From e897c05f77f99be29edc581834e324ab337fa432 Mon Sep 17 00:00:00 2001 From: shashvat-singham Date: Sun, 16 Aug 2026 03:06:15 +0530 Subject: [PATCH] Fix clip_by_global_norm producing NaN when g_norm == max_norm == 0 The trigger uses a strict comparison, so at equality the clip branch is selected and computes (t / g_norm) * max_norm. With all-zero updates and max_norm=0 that is (0 / 0) * 0 = NaN for every leaf: clipper = optax.clip_by_global_norm(0.0) updates, _ = clipper.update(jax.tree.map(jnp.zeros_like, grads), state) # every leaf is NaN Zero gradients are common (frozen parameters, saturated units), and a max_norm of exactly zero is reachable when the norm is driven by a schedule through inject_hyperparams, so the combination NaN-poisons an update mid-training rather than raising anywhere visible. Use a non-strict comparison instead. Clipping updates to exactly their own norm is the identity, so taking the pass-through branch at equality is mathematically unchanged (and exact, where the clip branch is only exact up to rounding) while avoiding the 0/0. unitwise_clip in the same module already guards its division with div_eps for exactly this reason; clip_by_global_norm was the remaining unguarded one. --- optax/transforms/_clipping.py | 6 +++++- optax/transforms/_clipping_test.py | 17 +++++++++++++++++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/optax/transforms/_clipping.py b/optax/transforms/_clipping.py index 824998e53..95cefbf27 100644 --- a/optax/transforms/_clipping.py +++ b/optax/transforms/_clipping.py @@ -99,7 +99,11 @@ def update_fn(updates, state, params=None): # once analyzed how it affects backprop through update (e.g. meta-gradients) # g_norm = jnp.maximum(max_norm, g_norm) # updates = jax.tree.map(lambda t: (t / g_norm) * max_norm, updates) - trigger = jnp.squeeze(g_norm < max_norm) + # Non-strict comparison: clipping updates to their own norm is the identity, + # and taking the pass-through branch at equality avoids computing 0/0 = NaN + # when the global norm and max_norm are both zero (e.g. all-zero gradients + # under a schedule-driven max_norm that reaches zero). + trigger = jnp.squeeze(g_norm <= max_norm) utils.check_rank(trigger, 0) # A scalar. def clip_fn(t): diff --git a/optax/transforms/_clipping_test.py b/optax/transforms/_clipping_test.py index d7fcdb45e..4b2f03879 100644 --- a/optax/transforms/_clipping_test.py +++ b/optax/transforms/_clipping_test.py @@ -75,6 +75,23 @@ def test_clip_by_global_norm(self): updates_step, _ = clipper.update(self.per_step_updates, None) test_utils.assert_trees_all_close(updates, updates_step) + def test_clip_by_global_norm_zero_norm_zero_max_norm(self): + # 0 / 0 in the clip branch used to produce NaN updates when the global + # norm and max_norm were both zero (e.g. all-zero gradients under a + # schedule-driven max_norm that reaches zero). + clipper = _clipping.clip_by_global_norm(0.0) + zero_updates = jax.tree.map(jnp.zeros_like, self.per_step_updates) + updates, _ = clipper.update(zero_updates, None) + test_utils.assert_trees_all_close(updates, zero_updates) + + def test_clip_by_global_norm_at_equality_is_identity(self): + # Clipping updates to exactly their own norm is a no-op: pins the + # boundary semantics of the g_norm <= max_norm comparison. + g_norm = optax.tree.norm(self.per_step_updates) + clipper = _clipping.clip_by_global_norm(g_norm) + updates, _ = clipper.update(self.per_step_updates, None) + test_utils.assert_trees_all_close(updates, self.per_step_updates) + def test_adaptive_grad_clip_with_axis(self): """Test adaptive_grad_clip with custom axis parameter."""