class_docstr: Abstract base class for boxed Metadata.
  
  ``AxisMetadata`` enables arbitrary, per axis metadata for variables.
  By using ``unbox`` the metadata is stripped away to obtain the original
  variables. By using unboxing, most code handling variables does not need
  to handle ``AxisMetadata`` specifically, but can directly operate on the JAX
  arrays that they wrap.
  
  Additionally, ``AxisMetadata`` supports updating metadata whenever an axis
  is added or removed by a functional transformation
  (e.g.: ``nn.scan`` or ``nn.vmap``) using the ``add_axis`` and ``remove_axis``
  methods.
  
  By extending ``AxisMetadata``, custom metadata can be stored. See
  ``Partitioned`` for a specific implementation.
docstr: Removes an axis from the axis metadata.
  
  Note that add_axis and remove_axis should act as each other's inverse
  (meaning: ``x.remove_axis(i, p).add_axis(i, p) == x``)
  
  Args:
    index: The position of the axis that is to be removed
    params: An arbitrary dictionary of parameters passed by the transformation
      that introduced the axis (e.g.: ``nn.scan`` or ``nn.vmap``). The
      user passes this dictionary as the `metadata_param` argument to the
      transformation.
  Returns:
    A new instance of the same type as self and with the same ``unbox``
    content with updated axis metadata.
docstr_tok: ["Removes","an","axis","from","the","axis","metadata",".","Note","that","add_axis","and","remove_axis","should","act","as","each","other","'s","inverse","(","meaning",":","``","x.remove_axis","(","i",",","p",")",".add_axis","(","i",",","p",")","==","x","``",")","Args",":","index",":","The","position","of","the","axis","that","is","to","be","removed","params",":","An","arbitrary","dictionary","of","parameters","passed","by","the","transformation","that","introduced","the","axis","(","e.g",".",":","``","nn.scan","``","or","``","nn.vmap","``",")",".","The","user","passes","this","dictionary","as","the","`","metadata_param","`","argument","to","the","transformation",".","Returns",":","A","new","instance","of","the","same","type","as","self","and","with","the","same","``","unbox","``","content","with","updated","axis","metadata","."]
code: def remove_axis(
      self: TAxisMetadata, index: int, params: Dict[Any, Any]
  ) -> TAxisMetadata:
      """Removes an axis from the axis metadata.
  
      Note that add_axis and remove_axis should act as each other's inverse
      (meaning: ``x.remove_axis(i, p).add_axis(i, p) == x``)
  
      Args:
        index: The position of the axis that is to be removed
        params: An arbitrary dictionary of parameters passed by the transformation
          that introduced the axis (e.g.: ``nn.scan`` or ``nn.vmap``). The
          user passes this dictionary as the `metadata_param` argument to the
          transformation.
      Returns:
        A new instance of the same type as self and with the same ``unbox``
        content with updated axis metadata.
      """
      pass
  
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 4 for calling the AxisMetadata obj's underlying member method code with example usage: obj.remove_axis(index, params) without return types