Right way to deal with zero gradient norm in jax.experimental.optimizers.clip_grads

[vuoristo]

The following snippet raises a FloatingPointError due to the gradient having zero norm. In my use case it's plausible that sometimes the gradient norm to some leaves would be zero. What's right way to deal with this?

import jax
import jax.numpy as jnp
from jax.experimental import optimizers

jax.config.update('jax_debug_nans', True)

optimizers.clip_grads(jnp.array([0]), 1)

[C-J-Cundy]

A super simple way to stop debug_nans firing an error, which shouldn't make any noticeable difference to the clipping as long as max_norm isn't really small:

from jax.tree_util import tree_map
from jax.experimental.optimizers import l2_norm
import jax.numpy as jnp
import jax

def safe_clip_grads(grad_tree, max_norm):
  """Clip gradients stored as a pytree of arrays to maximum norm `max_norm`."""
  norm = l2_norm(grad_tree)
  eps = 1e-9
  normalize = lambda g: jnp.where(norm < max_norm, g, g * max_norm / (norm + eps))
  return tree_map(normalize, grad_tree)


jax.config.update('jax_debug_nans', True)

safe_clip_grads(jnp.array([0]), 1)

[vuoristo]

Thanks! That's probably the right way to deal with it.