fix(transform): prevent NaN in Adam-family optimizers on float16 zero-gradient steps

[somuai]

Resolves #1754.

Problem

scale_by_adam, scale_by_amsgrad, scale_by_belief, scale_by_yogi, and scale_by_radam in optax/_src/transform.py compute parameter updates as m / (sqrt(v + eps_root) + eps). When parameters and gradients are in float16, Python float eps (default 1e-8, or 1e-16 for AdaBelief) underflows below the IEEE 754 float16 subnormal threshold (~5.96e-8) to 0.0.

On any step with an exact zero gradient (e.g. masked tokens, frozen layers, unused embedding rows), both moment estimates are 0.0, resulting in sqrt(0.0) + 0.0 == 0.0 in the denominator and producing 0.0 / 0.0 = NaN. Once generated, NaN values propagate through all subsequent training steps.

Solution

• Refactored update mapping across scale_by_adam, scale_by_amsgrad, scale_by_belief, scale_by_yogi, and _radam_update to evaluate the denominator in at least float32 precision using jnp.promote_types(v.dtype, jnp.float32).
• Applied division as (m / denom).astype(m.dtype), avoiding denominator underflow while cleanly preserving imaginary components for complex parameters (complex64).
• Added regression tests in optax/_src/transform_test.py:
  • test_adam_family_no_nan_on_float16_zero_grad: verifies all 5 optimizers produce exact 0.0 updates without NaNs on zero-gradient float16 inputs.
  • test_adam_family_complex_parameters: verifies complex parameter updates retain imaginary components.

[sylvesterkaczmarek]

The ArrayLike epsilon path can still reproduce this: eps=jnp.asarray(1e-8, dtype=jnp.float16) is already zero before it is added to the promoted denominator, so zero gradients can still produce 0/0. Promote/cast eps and eps_root to the safe denominator dtype (or reject non-representable epsilons) and add an array-epsilon regression test.

[somuai]

Great catch @sylvesterkaczmarek! Addressed in commit 603b694:

1. Promoted ArrayLike eps and eps_root: In _update across scale_by_adam, scale_by_amsgrad, scale_by_belief, scale_by_yogi, and _radam_update, cast both eps and eps_root to safe_dtype before evaluating the denominator (safe_eps = jnp.asarray(eps, dtype=safe_dtype)).
2. Rejected non-representable / non-positive eps: Added validation in each optimizer constructor (if jnp.any(jnp.asarray(eps) <= 0): raise ValueError(...)) to explicitly reject underflowed or zero/negative epsilons (such as eps=jnp.asarray(1e-8, dtype=jnp.float16) where 1e-8 underflows to 0.0).
3. Added array-epsilon regression test: Added test_adam_family_array_epsilon in optax/_src/transform_test.py verifying that valid ArrayLike epsilons (jnp.asarray(1e-4, dtype=jnp.float32)) compute finite updates on float16 zero-gradient steps without NaNs, and that underflowed float16 epsilons raise a ValueError.

Ready for another look!