Fix negative CTC loss for confident alignments in float32

[rajasekharporeddy]

Fixes #1771

Summary

In float32 arithmetic, optax.ctc_loss and optax.ctc_loss_with_forward_probs can round below zero ($< 0$) for highly confident alignments due to precision limits in log_softmax and logaddexp. This PR enforces the mathematical bounds by clamping forward log-probabilities from above at 0.0 and per-sequence loss from below at 0.0.

Changes

1. In optax.losses.ctc_loss_with_forward_probs:
   • Clamped logalpha_phi and logalpha_emit from above at 0.0 using jnp.minimum(..., 0.0).
   • Clamped per_seq_loss from below at 0.0 using jnp.maximum(per_seq_loss, 0.0).
2. In optax/losses/_classification_test.py:
   • Added unit test test_confident_alignment_nonnegative in CTCTest to verify that both ctc_loss and ctc_loss_with_forward_probs maintain non-negative loss and non-positive log-probabilities under confident alignments.