Fix argument order in scale_by_distance_over_gradients

The helper function _tx is defined with parameters (g, d, g_sos), but
jax.tree.map was calling it with (max_dist, g_sos, updates). This caused
_tx to interpret the accumulated gradient sum-of-squares as d and the
raw updates (which can be negative) as g_sos. When the raw gradient had
any negative entry, the code would execute jnp.sqrt(g_sos + eps) with a
negative argument and produce NaN.

This fix corrects the argument order to (updates, max_dist, g_sos) to
match the function signature.

Author: Aaryan549

optax/_src/transform.py

@@ -1406,7 +1406,7 @@ def _tx(g, d, g_sos):
       eta = global_scale * (d / jnp.sqrt(g_sos + eps))
       return eta * g
 
-    updates = jax.tree.map(_tx, max_dist, g_sos, updates)
+    updates = jax.tree.map(_tx, updates, max_dist, g_sos)
 
     # new state
     state = ScaleByDistanceOverGradientsState(