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