def concatenate(
    arrays: Sequence[NamedArrayBase], axis_name: AxisName
) -> NamedArray | NamedArrayView:
  """Concatenates a sequence of named arrays along a named axis.

  Args:
    arrays: Sequence of named arrays to concatenate, which must have the same
      shape except for the axis given by ``axis_name``.
    axis_name: Name of the axis along which to concatenate. Must be present in
      all arrays.

  Returns:
    A new concatenated array.
  """
  ndims = set(len(array.positional_shape) for array in arrays)
  if len(ndims) != 1:
    raise ValueError(
        "All arrays must have the same number of positional axes, but got"
        f" {ndims}"
    )
  (ndim,) = ndims

  orig_positional_axes = [TmpPosAxisMarker() for _ in range(ndim)]
  arrays_along_axis = [
      array.tag(*orig_positional_axes).untag(axis_name) for array in arrays
  ]
  concatenated = nmap(jnp.concatenate)(arrays_along_axis)
  return concatenated.tag(axis_name).untag(*orig_positional_axes)
Concatenates a sequence of named arrays along a named axis.

  Args:
    arrays: Sequence of named arrays to concatenate, which must have the same
      shape except for the axis given by ``axis_name``.
    axis_name: Name of the axis along which to concatenate. Must be present in
      all arrays.

  Returns:
    A new concatenated array.
  
penzai/core/named_axes.py