Fix : division by zero in sampling when temperature is 0.0

gemma/gm/text/_sampling.py

@@ -58,6 +58,8 @@ class RandomSampling(SamplingMethod):
 
   @typechecked
   def get_next_tokens(self, logits: Float['*B V'], rng: PRNGKey) -> Int['*B']:
+    if self.temperature < 1e-6:
+      return Greedy().get_next_tokens(logits, rng)
     return jax.random.categorical(rng, logits / self.temperature, axis=-1)
 
 
@@ -70,6 +72,9 @@ class TopkSampling(SamplingMethod):
 
   @typechecked
   def get_next_tokens(self, logits: Float['*B V'], rng: PRNGKey) -> Int['*B']:
+    if self.temperature < 1e-6:
+      return Greedy().get_next_tokens(logits, rng)
+
     logits, batch_shape = enp.flatten(logits, '... V')
 
     batch_size = logits.shape[0]
@@ -91,6 +96,9 @@ class TopPSampling(SamplingMethod):
 
   @typechecked
   def get_next_tokens(self, logits: Float['... V'], rng: PRNGKey) -> Int['...']:
+    if self.temperature < 1e-6:
+      return Greedy().get_next_tokens(logits, rng)
+
     # temperature scaling
     logits = logits / self.temperature
 
@@ -115,4 +123,3 @@ def get_next_tokens(self, logits: Float['... V'], rng: PRNGKey) -> Int['...']:
       )
 
     return jax.random.categorical(rng, logits, axis=-1)
-

gemma/gm/text/_sampling_test.py

@@ -79,3 +79,23 @@ def test_top1_sampling_matches_greedy_sampling():
   tokens_top1 = top1_sampling.get_next_tokens(logits, rng)
   np.testing.assert_array_equal(tokens_greedy, tokens_top1)
 
+
+
+def test_zero_temperature_behavior():
+  rng = jax.random.PRNGKey(0)
+  logits = jax.numpy.array([[10.0, 5.0]])
+
+  # Test RandomSampling
+  sampler = gm.text.RandomSampling(temperature=0.0)
+  tokens = sampler.get_next_tokens(logits, rng)
+  np.testing.assert_array_equal(tokens, [0])
+
+  # Test TopkSampling
+  sampler = gm.text.TopkSampling(k=5, temperature=0.0)
+  tokens = sampler.get_next_tokens(logits, rng)
+  np.testing.assert_array_equal(tokens, [0])
+
+  # Test TopPSampling
+  sampler = gm.text.TopPSampling(p=0.9, temperature=0.0)
+  tokens = sampler.get_next_tokens(logits, rng)
+  np.testing.assert_array_equal(tokens, [0])