Skip to content

fix: prevent division by zero in sampling when temperature is 0.0 (#562)#585

Open
irhyl wants to merge 1 commit intogoogle-deepmind:mainfrom
gsoc-2026:fix/562-zero-temp-sampling
Open

fix: prevent division by zero in sampling when temperature is 0.0 (#562)#585
irhyl wants to merge 1 commit intogoogle-deepmind:mainfrom
gsoc-2026:fix/562-zero-temp-sampling

Conversation

@irhyl
Copy link

@irhyl irhyl commented Feb 22, 2026

Summary

Fixes #562 — Setting temperature=0.0 in RandomSampling, TopkSampling, or
TopPSampling causes division by zero, producing inf/NaN logits and incorrect
token predictions.

Changes

gemma/gm/text/_sampling.py — Add a guard clause at the top of
get_next_tokens() in all three temperature-dependent sampling classes:

if self.temperature < 1e-6:
    return jnp.argmax(logits, axis=-1)

…ogle-deepmind#562)

Add guard clause in RandomSampling, TopkSampling, and TopPSampling to
fall back to greedy (argmax) when temperature < 1e-6, preventing inf/NaN
from division by zero.

Add tests to verify zero-temperature behavior matches Greedy sampling.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Bug: Division by zero in sampling methods when temperature is 0.0

1 participant