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

[jaideeppyne]

What

The custom JVP of weighted_logsoftmax (used by safe_softmax_cross_entropy) multiplies weights_dot by log_softmax(x). At a masked class where x_i = -inf, log_softmax(x)_i = -inf, so this term becomes weights_dot * -inf, producing NaN tangents — even for a zero tangent direction.

Why it's a bug

The primal already implements the 0 log 0 = 0 convention so masked -inf entries stay finite (per its docstring). The gradient rule did not, so jax.jvp(weighted_logsoftmax, primals, (0, 0)) and jax.jacfwd(safe_softmax_cross_entropy) return NaN on inputs with -inf logits.

import jax, jax.numpy as jnp, optax
logits = jnp.array([-jnp.inf, 0., 0.]); labels = jnp.array([0., 1., 0.])
_, t = jax.jvp(optax.losses.weighted_logsoftmax, (logits, labels),
               (jnp.zeros_like(logits), jnp.zeros_like(labels)))
print(t)  # before: [nan 0. 0.]   after: [0. 0. 0.]

Fix

Substitute a finite value for -inf entries of log_softmax(x) in the weights_dot term. This leaves the gradient unchanged wherever log_softmax is finite and removes the spurious NaN otherwise.

Tests

Adds a forward-mode regression test over inputs with -inf logits (fails before, passes after). All existing _classification tests pass; ruff clean.