[Aditya-B-007]

Prevent NaN updates in scale_by_rms/stddev on float16 zero-gradi…ent steps

eps=1e-8 underflows to 0.0 when evaluated in float16, making the denominator zero and producing inf * 0 = NaN on the first step with zero gradients. This affected scale_by_rms, scale_by_stddev, and therefore optax.rmsprop.

Fix:

• Compute the scaling denominator in float32 so eps remains 1e-8.
• Floor variance at jnp.finfo(float32).tiny before taking the square root, so a zero EMA can never produce an infinite scale factor.
• Cast the update product back to the gradient dtype, preserving float16.

Both eps_in_sqrt=True (rsqrt path) and eps_in_sqrt=False (1/sqrt+eps path) are fixed. bias_correction on/off is also covered.

[Aditya-B-007]

Fixed issue #1784 . Please do take a look

[Aditya-B-007]

Bug is fixed can be merged
Updated the fix to use a _safe_eps helper instead of casting update arrays.
_safe_eps(g, eps) clamps ϵ to jnp.finfo(real_dtype).tiny, ensuring that float16 denominators never underflow to 0.0 while preserving 1e-8 for standard float/complex dtypes. This resolves the float16 zero-gradient NaN issue without modifying or stripping the imaginary components of complex64 gradients.

[rajasekharporeddy]

Thanks for working on this and contributing a fix for #1784, @Aditya-B-007!

We've addressed the float16 eps / eps_root underflow bugs repo-wide in #1835. Rather than silently promoting the denominator to float32, Optax now validates eps and eps_root at trace time via numerics.maybe_static_check_eps and raises an explicit ValueError when 0 < eps < jnp.finfo(dtype).tiny, guiding users to either set a larger eps (e.g., eps=1e-4 for float16) or use bfloat16/float32 optimizer state.

Closing this PR as superseded by #1835. Thanks again for your effort!

[Aditya-B-007]

Thank you.