fix: prevent division by zero in sampling when temperature is 0.0

[ChiragTrivedi06]

Fix: Numerical Stability Guard for Sampling Methods

Problem

Currently, the RandomSampling, TopkSampling, and TopPSampling classes perform a direct division by the temperature parameter. When temperature=0.0, this leads to a division-by-zero error, causing JAX/XLA to produce NaN or Inf logits and crashing the sampling pipeline.

Proposed Changes

Instead of a hard switch to a separate Greedy implementation, this PR introduces a threshold-based guard within the sampling classes. This ensures numerical stability while maintaining a unified API for consumers.

Key Implementation Details:

• Threshold Guard: Added scaled_logits = logits if self.temperature < 1e-6 else logits / self.temperature across all sampling methods.
• Unified Interface: Avoids forced object swapping; users can now toggle between stochastic and deterministic decoding simply by adjusting the temperature hyperparameter.
• Top-P/Top-K Compatibility: Ensures that filtering operations (sorting, cumulative distribution) remain stable even at the limit of zero temperature.

---

Technical Rationale

1. Mathematical Continuity: As $\tau \to 0$, categorical sampling naturally converges to a greedy selection. This change ensures the implementation reflects this mathematical property without hitting hardware-level numerical exceptions.
2. Infrastructure Compatibility: Many evaluation and training loops are designed around a single sampling configuration. Forcing a class change (to Greedy) for $\tau=0$ introduces unnecessary branching logic in downstream code.

Verification

• Numerical Stability: Verified that $\tau=0.0$ no longer produces NaN logits or XLA errors.
• Behavioral Consistency: Confirmed that low-temperature sampling correctly priorities the most probable tokens (greedy behavior).
• Cross-Backend Testing: Verified consistent output across CPU, GPU, and TPU.
• Unit Testing: Updated _sampling_test.py with cases for near-zero and zero temperature.

Closes #562

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.