code: def remove_axis(
      state: State, index: int, params: tp.Mapping[tp.Any, tp.Any]
  ) -> State:
      axis_name = _get_partition_name(params)
  
      def _remove_axis(x: tp.Any):
          if isinstance(x, variables.Variable):
              if isinstance(x, HasSharding) and x.sharding is not None:
                  sharding = list(x.sharding)
                  assert sharding.pop(index) == axis_name
                  x.sharding = tuple(sharding)
              x.remove_axis(axis_name, index)
          return x
  
      return jax.tree_map(
          _remove_axis,
          state,
          is_leaf=lambda x: isinstance(x, variables.Variable),
      )
  
imports: import functools
  import typing
  import jax
  from jax.experimental import maps
  from jax.sharding import Mesh, PartitionSpec
  from flax.experimental.nnx.nnx import variables
  from flax.experimental.nnx.nnx.pytreelib import TreeNode
  from flax.experimental.nnx.nnx.state import State
label_desc: Use image node_id 2 for calling a global function with example usage: remove_axis(state, index, params) and returns: jax, x