code: def clip_grads(grad_tree, max_norm):
      """Clip gradients stored as a pytree of arrays to maximum norm `max_norm`."""
      norm = l2_norm(grad_tree)
      normalize = lambda g: jnp.where(
          norm < max_norm, g, g * (max_norm / norm)
      )
      return tree_map(normalize, grad_tree)
  
imports: from __future__ import annotations
  from typing import Any, Callable, NamedTuple
  from collections import namedtuple
  import functools
  from functools import partial
  import jax.numpy
  from jax._src.util import safe_zip, safe_map, unzip2
  from jax import tree_util
  from jax.tree_util import tree_map, tree_flatten, tree_unflatten, register_pytree_node
label_desc: Use image node_id 18 for calling a global function with example usage: clip_grads(grad_tree, max_norm) and returns: tree_map