fix(losses): avoid NaN gradient in weighted_logsoftmax at -inf logits

optax/losses/_classification.py

@@ -198,11 +198,19 @@ def _weighted_logsoftmax_jvp(primals, tangents):
   result = jnp.where(
       weights != 0.0, weights * logsoftmax_x, jnp.zeros_like(logsoftmax_x)
   )
+  # Apply the same ``0 * log(0) = 0`` convention as the primal to the
+  # ``weights_dot`` term. Entries where ``logsoftmax_x`` is ``-inf`` (e.g.
+  # ``x_i = -inf`` at a masked class) otherwise yield ``weights_dot * -inf``,
+  # which is ``nan`` even for a zero tangent direction. Substituting a finite
+  # value keeps the derivative correct wherever ``logsoftmax_x`` is finite.
+  safe_logsoftmax_x = jnp.where(
+      jnp.isneginf(logsoftmax_x), jnp.zeros_like(logsoftmax_x), logsoftmax_x
+  )
   out_tangents = (
       weights * x_dot
       - weights
       * jnp.sum(x_dot * jax.nn.softmax(x, axis=-1), axis=-1, keepdims=True)
-      + weights_dot * logsoftmax_x
+      + weights_dot * safe_logsoftmax_x
   )
   return result, out_tangents
 

optax/losses/_classification_test.py

@@ -163,6 +163,25 @@ def test_gradient(self):
         order=1,
     )
 
+  def test_forward_mode_gradient_finite_with_infinite_logits(self):
+    """Forward-mode differentiation must respect the ``0 log 0 = 0`` rule.
+
+    Rows with ``-inf`` logits at masked (zero-weight) classes previously made
+    the custom JVP of ``weighted_logsoftmax`` compute ``weights_dot * -inf``,
+    yielding ``nan`` tangents even for a zero tangent direction.
+    """
+    logits, labels = self.ys, self.ts
+    # A zero tangent direction must produce a finite (zero) output tangent.
+    _, out_tangents = jax.jvp(
+        _classification.weighted_logsoftmax,
+        (logits, labels),
+        (jnp.zeros_like(logits), jnp.zeros_like(labels)),
+    )
+    np.testing.assert_array_equal(out_tangents, np.zeros_like(out_tangents))
+    # Forward-mode Jacobian w.r.t. the logits must be free of NaNs.
+    jac = jax.jacfwd(_classification.safe_softmax_cross_entropy)(logits, labels)
+    self.assertFalse(np.any(np.isnan(jac)))
+
   def test_against_plain_implementation(self):
     """Tests against plain implementation which does not handle -inf."""