[copybara-service]

Raise a ValueError at trace time when eps underflows in the optimizer state dtype

Positive eps values smaller than the smallest normal number of the dtype they are added in (jnp.finfo(dtype).tiny) are flushed to zero. For example the default eps=1e-8 is flushed to zero in float16 (smallest normal 2^-14 ~= 6.1e-5), which turns zero-gradient steps into 0 / 0 = NaN (#1784, #1789).

Rather than silently changing the requested eps, this change:

• Adds numerics.maybe_static_check_eps(eps, tree, name), which raises a ValueError when a concrete eps satisfies 0 < eps < jnp.finfo(dtype).tiny for the dtype of eps + leaf. Traced eps (e.g. from inject_hyperparams) and eps=0 are not checked.
• Calls it on eps (and eps_root) in every transform that adds eps to a per-parameter gradient-statistics accumulator:
  • optax/_src/transform.py: scale_by_rss, scale_by_rms, scale_by_stddev, scale_by_adam, scale_by_amsgrad, scale_by_adamax, scale_by_adadelta, scale_by_adan, scale_by_belief, scale_by_yogi, scale_by_radam, scale_by_sm3, scale_by_novograd. optax.sm3 gains an eps argument so the check is actionable.
  • optax/_src/factorized.py: scale_by_factored_rms (epsilon, used by adafactor).
  • optax/contrib: acprop, ademamix, simplified_ademamix, adopt, dadapt_adamw, madgrad, momo_adam, sophia.
• Default eps values are unchanged. As a consequence, optimizers using these transforms with default eps now raise for float16 optimizer state; pass a larger eps (e.g. eps=1e-4) or keep the optimizer state in bfloat16/float32. float16 is removed from the dtype-stability test matrices and the sharding tests use bfloat16.

Thanks to @vishesh9131 (#1785) and @Anu27n (#1790) for reporting and the initial fixes.

[vishesh9131]

NP