code: def _norm_except_axis(
      v: Array, norm: Callable[[Array], Scalar], axis: Optional[int]
  ):
      if axis is None:
          return norm(v)
      else:
          return jax.vmap(norm, in_axes=axis, out_axes=axis)(v)
  
file_path: equinox/equinox/nn/_weight_norm.py
imports: import functools
  from collections.abc import Callable
  from typing import Generic, Optional, TypeVar
  import jax
  import jax.numpy
  from jaxtyping import Array, PRNGKeyArray, Scalar
  from .._module import field, Module
  from .._tree import tree_at
label_desc: Use image node_id 1 for calling a global function with example usage: _norm_except_axis(v, norm, axis) and returns: norm, jax