def _l2_norm(self, flat_list):
    """Returns an L2-like norm of all elements of all tensors in `flat_list`.

    Args:
      flat_list: Collection of tensors as a flat list (rather than, e.g., a
          tree).

    Returns:
      A scalar value computed as if all the tensors in `flat_list` were joined
      and flattened into a single vector, and then the L2 norm of that vector
      was calculated.
    """
    if fastmath.is_backend(fastmath.Backend.JAX):
      norm = jnp.sqrt(sum(jnp.vdot(x, x) for x in flat_list))
    else:  # TODO(lukaszkaiser): add vdot to TF-numpy
      norm = jnp.sqrt(sum(jnp.sum(x*x) for x in flat_list))
    return norm
Returns an L2-like norm of all elements of all tensors in `flat_list`.

    Args:
      flat_list: Collection of tensors as a flat list (rather than, e.g., a
          tree).

    Returns:
      A scalar value computed as if all the tensors in `flat_list` were joined
      and flattened into a single vector, and then the L2 norm of that vector
      was calculated.
    
trax/optimizers/base.py