def unstack(
    array: NamedArrayBase, axis_name: AxisName
) -> Sequence[NamedArray | NamedArrayView]:
  """Splits a named array across a given named axis.

  Args:
    array: The array to split.
    axis_name: Name of the axis along which to unstack.

  Returns:
    A sequence of slices across the given axis.
  """
  return [array[{axis_name: i}] for i in range(array.named_shape[axis_name])]
Splits a named array across a given named axis.

  Args:
    array: The array to split.
    axis_name: Name of the axis along which to unstack.

  Returns:
    A sequence of slices across the given axis.
  
penzai/core/named_axes.py