code: def map_axis_meta(
      fn: Callable[[AxisMetadata[Any]], Any], tree: Any
  ) -> Any:
      """Maps over all PyTree nodes that are AxisMetadata instances."""
  
      def wrapper(x):
          if isinstance(x, AxisMetadata):
              return fn(x)
          else:
              return x
  
      return jax.tree_map(wrapper, tree, is_leaf=is_axis_metadata)
  
imports: import abc
  import functools
  from typing import Any, Callable, Dict, Generic, Optional, Tuple, TypeVar, Union
  import jax
  from jax.experimental import maps
  from flax import errors, struct
label_desc: Use image node_id 2 for calling a global function with example usage: map_axis_meta(fn, tree) and returns: jax, fn, x