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!