arange

Contents

arange#

penzai.core.named_axes.arange(name: str, start: int, stop: int | None = None, step: int | None = None, dtype: jax.typing.DTypeLike | None = None) → NamedArray[source]#

Convenience function to create a range along a named axis.

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

Parameters:

  • 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.

Contents