jax-ml/jax · error · ValueError

axis is out of range.

Error message

axis is out of range.

What it means

roll in Mosaic validates that the roll axis is a valid non-negative index into x's shape (negative indices are not normalized like jnp.roll). Passing axis < 0 or axis >= x.ndim raises this ValueError immediately.

Source

Thrown at jax/_src/pallas/mosaic/primitives.py:131

batching.primitive_batchers[bitcast_p] = _bitcast_batch_rule


roll_p = jax_core.Primitive("roll")


def roll(
    x: jax.Array,
    shift: jax.Array | int,
    axis: int,
    *,
    stride: int | None = None,
    stride_axis: int | None = None,
) -> jax.Array:
  if isinstance(shift, int) and shift < 0:
    raise ValueError("shift must be non-negative.")
  if axis < 0 or axis >= len(x.shape):
    raise ValueError("axis is out of range.")
  if (stride is None) != (stride_axis is None):
    raise ValueError("stride and stride_axis must be both specified or not.")
  if stride is not None and stride_axis is not None:
    if stride < 0:
      raise ValueError("stride must be non-negative.")
    if stride_axis < 0 or stride_axis >= len(x.shape):
      raise ValueError("stride_axis is out of range")
    if axis == stride_axis:
      raise ValueError("expected axis and stride_axis are different.")
  return roll_p.bind(
      x, shift, axis=axis, stride=stride, stride_axis=stride_axis
  )


@roll_p.def_abstract_eval
def _roll_abstract_eval(x, shift, **_):
  del shift
  return x

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Normalize negative axes first: axis = axis % len(x.shape)
2. Assert axis is in range before calling roll
3. Use the last-axis spelling explicitly, e.g. axis=x.ndim - 1 instead of -1

Example fix

# before
y = roll(x, 2, axis=-1)

# after
y = roll(x, 2, axis=len(x.shape) - 1)

assert 0 <= axis < len(x.shape), f"axis {axis} out of range for shape {x.shape}"
# normalize negative axes like jnp:
axis = axis % len(x.shape)

When it happens

Trigger: Calling mosaic roll with a negative axis (roll(x, 2, axis=-1)) or an axis index equal to or beyond the rank of x.

Common situations: Translating jnp.roll or np.roll calls that idiomatically use axis=-1/-2 into a Pallas kernel; changing tensor rank (e.g. adding a batch dim) so a previously valid axis constant is now out of range.

Related errors

• stride_axis is out of range
• Not implemented: bitcast 1D
• Not implemented: the 2nd minor dim can not be perfectly pack
• shift must be non-negative.
• stride and stride_axis must be both specified or not.

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/21c145cd69fcf1ad. Report an issue: GitHub.