jax-ml/jax · error · ValueError

expected axis and stride_axis are different.

Error message

expected axis and stride_axis are different.

What it means

Strided roll rolls along `axis` while stepping along `stride_axis`; these must be different dimensions for the operation to be meaningful. Passing the same axis for both is rejected because the stride would fold into the roll itself.

Source

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

    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


def _roll_lowering_rule(
    ctx: mlir.LoweringRuleContext, x, shift, *, axis, stride, stride_axis
):
  def _roll(x, shift):
    if stride is None:
      return jnp.roll(x, shift, axis)
    outputs = [

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Choose a stride_axis different from axis (typically the dimension you are iterating blocks over)
2. If you only want a plain roll, drop stride and stride_axis entirely
3. Add an assert axis != stride_axis in your kernel wrapper for early failure

Example fix

# before
y = roll(x, 2, axis=1, stride=4, stride_axis=1)

# after
y = roll(x, 2, axis=1, stride=4, stride_axis=0)
# or plain roll:
y = roll(x, 2, axis=1)

if stride is not None:
  assert axis != stride_axis, "axis and stride_axis must differ"
# drop strided args entirely when unused:
if stride is None:
  y = roll(x, shift, axis=axis)
else:
  y = roll(x, shift, axis=axis, stride=stride, stride_axis=stride_axis)

When it happens

Trigger: roll(x, shift, axis=1, stride=4, stride_axis=1).

Common situations: Defaulting both parameters to the same value in a wrapper; refactoring where stride_axis was copied from axis and never changed; misunderstanding the two-axis strided rolling API.

Related errors

• shift must be non-negative.
• stride and stride_axis must be both specified or not.
• stride must be non-negative.
• axis is out of range.
• stride_axis is out of range

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