def arange(
    name: str,
    start: int,
    stop: int | None = None,
    step: int | None = None,
    dtype: jax.typing.DTypeLike | None = None,
) -> NamedArray:
  """Convenience function to create a range along a named axis.

  This is shorthand for ``wrap(jnp.arange(...)).tag(name)``.

  Args:
    name: Name for the resulting axis.
    start: Start of the range. If ``stop`` is not provided, this is instead
      interpreted as ``stop``, with ``start`` implicitly set to 0, following
      `jnp.arange`.
    stop: End of the range.
    step: Step size (defaults to 1).
    dtype: Optional dtype for the result.

  Returns:
    `NamedArray` with a single named axis containing a range of integers.
  """
  return wrap(jnp.arange(start, stop, step, dtype)).tag(name)
Convenience function to create a range along a named axis.

  This is shorthand for ``wrap(jnp.arange(...)).tag(name)``.

  Args:
    name: Name for the resulting axis.
    start: Start of the range. If ``stop`` is not provided, this is instead
      interpreted as ``stop``, with ``start`` implicitly set to 0, following
      `jnp.arange`.
    stop: End of the range.
    step: Step size (defaults to 1).
    dtype: Optional dtype for the result.

  Returns:
    `NamedArray` with a single named axis containing a range of integers.
  
penzai/core/named_axes.py