jax-ml/jax · error · ValueError

stride and stride_axis must be both specified or not.

Error message

stride and stride_axis must be both specified or not.

What it means

Mosaic roll supports an optional strided mode, but the stride and stride_axis keyword arguments are coupled: they describe one stride specification and must be supplied together or both omitted. Supplying only one of them is contradictory and rejected.

Source

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

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. Pass both stride and stride_axis together, or pass neither
2. Audit helper wrappers that build kwargs dynamically to ensure the two are set jointly
3. If you do not need strided rolling, remove the stray argument

Example fix

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

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

assert (stride is None) == (stride_axis is None), "pass both or neither"
kwargs = {} if stride is None else dict(stride=stride, stride_axis=stride_axis)
y = roll(x, shift, axis=axis, **kwargs)

When it happens

Trigger: roll(x, shift, axis) with stride=4 but no stride_axis, or stride_axis=1 but no stride.

Common situations: Copy-pasting partial kwargs between roll call sites; refactoring where a stride parameter is threaded through but stride_axis is dropped by mistake; optional-argument plumbing with None defaults where one gets set conditionally.

Related errors

• shift must be non-negative.
• stride must be non-negative.
• 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/5ec4a779e998f49f. Report an issue: GitHub.