code: def _add_axis_to_metadata(
      fn, axis_pos, axis_name, axis_col="params_axes"
  ):
      """Insert a named axis to axes metadata."""
      # Handle In() / Out() scan axis marker types.
      if hasattr(axis_pos, "axis"):
          axis_pos = axis_pos.axis
  
      def insert_fn_leaf(names):
          if names is None:
              return names
          names = list(names)
          names.insert(axis_pos, axis_name)
          return tuple(names)
  
      def insert_fn(x):
          new_names = jax.tree_util.tree_map(
              insert_fn_leaf, x.names, is_leaf=_is_logical_spec
          )
          return x.replace(names=new_names)
  
      def remove_fn_leaf(names):
          if names is None:
              return names
          names = list(names)
          if names[axis_pos] != axis_name:
              raise ValueError(
                  f"Expected axis {axis_name} at position {axis_pos} in "
                  f"axis metadata {names}."
              )
          names.pop(axis_pos)
          return tuple(names)
  
      def remove_fn(x):
          new_names = jax.tree_util.tree_map(
              remove_fn_leaf, x.names, is_leaf=_is_logical_spec
          )
          return x.replace(names=new_names)
  
      return nn.transforms.map_variables(
          fn,
          axis_col,
          mutable=_is_mutable(axis_col),
          trans_in_fn=lambda tree: _tree_map_axes(remove_fn, tree),
          trans_out_fn=lambda tree: _tree_map_axes(insert_fn, tree),
      )
  
imports: import functools
  import re
  from typing import Any, Callable, Mapping, Optional, Tuple
  import flax
  from flax import linen
  from flax import struct
  from flax.core.frozen_dict import freeze
  from flax.core.frozen_dict import unfreeze
  from flax.core.lift import In
  from flax.core.lift import Out
  from flax.linen.spmd import _axis_rules
  from flax.linen.spmd import _AxisRules
  from flax.linen.spmd import _is_logical_spec
  from flax.linen.spmd import _with_sharding_constraint
  from flax.linen.spmd import Array
  from flax.linen.spmd import ArrayPytree
  from flax.linen.spmd import get_logical_axis_rules
  from flax.linen.spmd import logical_axis_rules
  from flax.linen.spmd import logical_to_mesh
  from flax.linen.spmd import logical_to_mesh_axes
  from flax.linen.spmd import LogicalPartitionSpec
  from flax.linen.spmd import LogicalPartitionSpecPytree
  from flax.linen.spmd import LogicalRules
  from flax.linen.spmd import PartitionSpecPytree
  from flax.linen.spmd import RulesFallback
  from flax.linen.spmd import set_logical_axis_rules
  from flax.linen.spmd import with_logical_constraint
  from flax.traverse_util import flatten_dict
  from flax.traverse_util import unflatten_dict
  import jax
label_desc: Use image node_id 8 for calling a global function with example usage: _add_axis_to_metadata(fn, axis_pos, axis_name, axis_col) and returns: nn, tuple, x, tuple, x, names, names