Add HOWTO: Gradient clipping

We currently don’t have a canonical implementation in Flax. Some optimizers have it built in.

Possible implementation:

def clip_grad_norm(grad, max_norm):
  norm = jnp.linalg.norm(jax.tree_util.tree_leaves(jax.tree_map(jnp.linalg.norm, grad)))
  factor = jnp.minimum(max_norm, max_norm / (norm + 1e-6))
  return jax.tree_map((lambda x: x * factor), grad)

[EyalRozenberg1]

Yes. I actually just saw that there is an implementation there for what we discussed. optax.clip_by_global_norm()

[jheek]

Shouldn’t this be part of optax?