def _partition_dict_index(self):
    """Helper to partition the indices of a dict-style index."""
    indexer = self.indexer
    assert isinstance(indexer, dict)
    # Our strategy will be to convert the indexed axis names into
    # positional axes, with
    # - named axes that are being introduced coming first,
    # - named axes that are being sliced (and thus preserved) coming next,
    # - named axes that are being removed or indexed using advanced indexing
    #   coming after,
    # - and finally, the positional axes of the original array.
    # We can then use ordinary positional indexing logic to retrieve the
    # result.
    sliced_axes = []
    introduced_axes = []
    indexed_axes = []

    for name, index in self.indexer.items():
      if index is None:
        # New axis.
        introduced_axes.append(name)
      elif isinstance(index, slice):
        # Sliced axis.
        sliced_axes.append(name)
      elif isinstance(index, int) or (
          isinstance(index, jax.Array | np.ndarray | NamedArrayBase)
          and jnp.issubdtype(index.dtype, np.integer)
      ):
        # Indexed axis.
        indexed_axes.append(name)
      else:
        raise TypeError(
            "Unsupported index for a named axis using dict-style index:"
            " expected a slice, an integer, an integer array, or None, but"
            f" got {index}"
        )

    input_prefix_order = (*sliced_axes, *indexed_axes)
    slice_order = (*introduced_axes, *sliced_axes, *indexed_axes)
    output_prefix_order = (*introduced_axes, *sliced_axes)

    return input_prefix_order, slice_order, output_prefix_order
Helper to partition the indices of a dict-style index.
penzai/core/named_axes.py