Fix clip_by_global_norm producing NaN when the global norm and max_norm are both zero

[shashvat-singham]

What

clip_by_global_norm selects between pass-through and clipping with a strict comparison:

trigger = jnp.squeeze(g_norm < max_norm)
...
return jax.lax.select(trigger, t, (t / g_norm.astype(t.dtype)) * max_norm)

At g_norm == max_norm the clip branch is selected. When both are zero, that branch computes (0 / 0) * 0 and every leaf of the update tree becomes NaN:

>>> clipper = optax.clip_by_global_norm(0.0)
>>> zero_grads = {"w": jnp.zeros(3)}
>>> clipper.update(zero_grads, clipper.init(zero_grads))[0]
{'w': Array([nan, nan, nan], dtype=float32)}

The neighbouring cases are all fine — zero grads under a positive max_norm pass through as zeros, and real grads under max_norm=0 clip to zeros — so this is exactly the equality corner.

Why it matters

Zero gradients are ordinary (frozen parameters, saturated units), and max_norm == 0.0 is reachable without anyone writing a literal zero: a schedule-driven max_norm through inject_hyperparams that decays to zero hits this the moment gradients are zero. The failure mode is silent NaN-poisoning of an update mid-training rather than an error anywhere visible.

Fix

Make the comparison non-strict (<=). Clipping updates to exactly their own norm is the identity, so taking the pass-through branch at equality is mathematically unchanged — and strictly better numerically, since the pass-through branch returns t exactly while the clip branch only recovers t up to rounding from (t / g_norm) * max_norm.

For precedent within the same module: unitwise_clip already guards its division with div_eps ("just prevents division by zero"); clip_by_global_norm was the remaining unguarded division.

I kept the existing lax.select structure untouched (the TODO about backprop-through-update / meta-gradients), since select's VJP routes cotangents through the selected branch only — at equality the derivative now comes from the identity branch, which is the mathematically sensible one-sided choice.

Tests

• test_clip_by_global_norm_zero_norm_zero_max_norm — fails on main (NaN), passes with the fix.
• test_clip_by_global_norm_at_equality_is_identity — pins the boundary semantics.

$ python -m pytest optax/transforms/_clipping_test.py -q
11 passed, 2 subtests passed

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

[shashvat-singham]

@googlebot I signed it!

[shashvat-singham]

For reviewers, on the two red checks:

• pyrefly: the 10 errors are all in _linesearch.py (lines 279/429 — dtype attribute and complex comparisons), which this PR does not touch. The same workflow is failing on main's last three runs with the same errors, so it's pre-existing on main rather than introduced here. Lint (ruff/flake8/pylint), doctests, and pre-commit all pass on this PR.
• cla/google: signed — re-triggered the check above.

[shashvat-singham]

@googlebot I signed it!

[shashvat-singham]

Update — the CLA is now signed and green (the earlier failure was a stale check from before signing; re-pushing the commit re-ran it).

Current state on e897c05: everything passes except pyrefly, and that failure is pre-existing on main, not from this PR:

• The reported errors are all in _linesearch.py (dtype on line 279, the complex/bool comparisons on line 429, and two bad-argument-types). The word clipping doesn't appear anywhere in the log.
• The same workflow is red on main's recent runs with the same errors.

Everything this PR does touch is green, including the full test matrix:

Pytest 3.10 ubuntu jax=0.5.3    pass
Pytest 3.10 ubuntu jax=newest   pass
Pytest 3.12 ubuntu jax=newest   pass
Pytest 3.12 ubuntu jax=nightly  pass
Pytest 3.13 ubuntu jax=newest   pass
ruff / flake8 / pylint / pre-commit / doctests / cla   pass

Happy to rebase if main gets the pyrefly issues fixed and you'd like a fully green run.