Fix NaN in scale_by_adam on float16 zero-gradient steps

[kartik1pandey]

Summary

Fixes NaN in scale_by_adam when used with float16 params on a zero-gradient step.

Changes

• scale_by_adam's update now adds eps/eps_root in at least float32 precision (jnp.promote_types(v.dtype, jnp.float32)) before casting the result back to the original dtype.
• Added test_adam_no_nan_on_float16_zero_grad regression test (float16 params, 3 steps of all-zero gradients, asserts every leaf stays finite).

Why

eps (a plain Python float, default 1e-8) is silently rounded to 0.0 when added to a float16 array, since float16's smallest representable subnormal is ~6e-8:

>>> jnp.asarray(1e-8, dtype=jnp.float16)
Array(0., dtype=float16)

On a step with an exactly-zero gradient (masked tokens, a frozen-then-unfrozen layer, etc.), both moment estimates are 0, so the sqrt(v) + eps safety denominator collapses to 0, producing 0 / 0 = NaN:

params = {"w": jnp.array([1.0, -2.0, 0.5], dtype=jnp.float16)}
grads = {"w": jnp.zeros_like(params["w"])}
opt = optax.adam(1e-3)
state = opt.init(params)
updates, state = opt.update(grads, state, params)
print(updates["w"])  # [nan nan nan]

Since NaN propagates through every subsequent step via apply_updates, this silently poisons the rest of training. The fix is a no-op for float32/float64 (verified bit-for-bit identical output before/after); only float16/bfloat16 behavior changes.

Fixes #1754.

Scoped to scale_by_adam only — scale_by_amsgrad, scale_by_belief, and scale_by_yogi share the same pattern (tracked in the linked issue); happy to follow up separately.

Testing

$ python -m pytest optax/_src/transform_test.py optax/_src/alias_test.py -q
560 passed, 62 skipped

ruff check is clean on both changed files.

[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.

[kartik1pandey]

@googlebot I signed it!

[winklemad]

The float32-promotion approach here looks right to me — computing the denominator in at least float32 keeps eps from rounding to 0.0 under float16 weak-type promotion, which is the actual cause.

One thing worth folding in before this merges: the same ... + eps denominator (the + eps outside the sqrt) is copy-pasted across several sibling optimizers in transform.py, and they underflow the same way, so a scale_by_adam-only fix leaves them NaN on the same input. I reproduced it against released optax 0.2.8:

import jax.numpy as jnp, optax
p = jnp.zeros((3,), jnp.float16); g = jnp.zeros((3,), jnp.float16)
for name, tx in [
    ('adam',    optax.scale_by_adam()),
    ('amsgrad', optax.scale_by_amsgrad()),
    ('belief',  optax.scale_by_belief()),
    ('rms',     optax.scale_by_rms(eps_in_sqrt=False)),
    ('stddev',  optax.scale_by_stddev(eps_in_sqrt=False)),
]:
    st = tx.init(p); u, _ = tx.update(g, st, p)
    print(name, bool(jnp.isnan(u).any()))
# adam/amsgrad/belief/rms/stddev -> all True

The matching denominators are:

• scale_by_amsgrad — m / (jnp.sqrt(v + eps_root) + eps) (same line as adam's)
• scale_by_belief — m / (jnp.sqrt(v) + eps)
• scale_by_rms / scale_by_stddev — 1 / (jnp.sqrt(...) + eps) on the eps_in_sqrt=False branch

Since it's the same fix each time, it might be cleanest to pull the safe-dtype denominator into a small shared helper and use it at each site, rather than patching adam alone — and then the test could parametrize over the affected optimizers. Happy to help with that if useful. Nice catch on the root cause either way.

[kartik1pandey]

Pulled the safe-dtype denominator into a shared
helper, numerics.add_eps_in_safe_dtype, and applied it to scale_by_amsgrad,
scale_by_belief, and both branches of scale_by_rms/scale_by_stddev.

One addition beyond your repro: the eps_in_sqrt=True branch NaNs too
(rsqrt(0) = inf, then inf * 0 = NaN) fixed the same way, keeping the
lax.rsqrt fusion so float32/float64 stay bit-for-bit unchanged.
scale_by_rms/scale_by_stddev also needed an explicit downcast after the
multiply so the fix doesn't silently upcast float16 updates to float32.

Verified: repro now passes on all 7 cases, bit-for-bit identical
output on float32/float64 vs. pre-fix, 653 passed/62 skipped on
transform/alias/numerics tests, ruff clean. Test is now parametrized over
all 7 cases per your suggestion.

Note: scale_by_yogi doesn't actually need this its default eps=1e-3
is too large to underflow in float16. scale_by_adamax does NaN on the
same input but via a different code path (no shared sqrt(...)+eps
denominator) flagging as a separate issue rather than folding in here.

[winklemad]

Nice — this is a cleaner home for it than my per-call change. Folding it into numerics.add_eps_in_safe_dtype and covering scale_by_amsgrad / scale_by_belief / both branches of scale_by_rms / scale_by_stddev in one place is the right call, and good catch on the eps_in_sqrt=True branch too. LGTM from my side — happy to give it another look once CI is green.

[kartik1pandey]

CI's green now all checks passing.
Ready whenever you get a chance to take another look. Thanks for the review!