jax-ml/jax · error · ValueError

stride must be non-negative.

Error message

stride must be non-negative.

What it means

When Mosaic roll is used in strided mode, the stride magnitude must be a non-negative integer. Negative strides are not supported by the underlying hardware lowering, so the public wrapper rejects them upfront.

Source

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

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


def _roll_lowering_rule(
    ctx: mlir.LoweringRuleContext, x, shift, *, axis, stride, stride_axis
):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use a non-negative stride; express direction via the axis/shift combination instead
2. Take abs(stride) if magnitude is what matters
3. Validate stride >= 0 at the call site with a clear error

Example fix

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

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

stride = abs(stride)
assert stride >= 0

When it happens

Trigger: roll(x, shift, axis, stride=-2, stride_axis=0) — any negative stride value together with a valid stride_axis.

Common situations: Deriving stride from a signed difference (e.g. element_size differences or direction flags) that can be negative; porting layout code where negative strides meant reversed traversal.

Related errors

• shift must be non-negative.
• stride and stride_axis must be both specified or not.
• expected axis and stride_axis are different.
• 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/bc94ac8bbdf5c591. Report an issue: GitHub.