jax-ml/jax · error · ValueError

shift must be non-negative.

Error message

shift must be non-negative.

What it means

The Mosaic roll primitive only accepts non-negative integer shift values; a negative shift would require reverse-direction hardware support that the primitive does not implement. The check runs on Python int shifts before binding the primitive.

Source

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

def _bitcast_batch_rule(batched_args, batch_axes, *, ty):
  return bitcast(*batched_args, ty=ty), batch_axes[0]

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, **_):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Convert negative shift to its positive equivalent: shift % x.shape[axis]
2. Clamp or validate shift before calling roll
3. Compute shift = abs(shift) with direction handled by which end you read from, if applicable

Example fix

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

# after
y = roll(x, shift=(-3) % x.shape[1], axis=1)

shift = shift % x.shape[axis] if isinstance(shift, int) else shift
assert not (isinstance(shift, int) and shift < 0)
y = roll(x, shift, axis=axis)

When it happens

Trigger: Calling mosaic roll with a literal negative int shift, e.g. roll(x, shift=-4, axis=1), including indirectly when a helper (like _roll) forwards a user-supplied constant.

Common situations: Porting numpy/jnp.roll code that uses negative shifts (which conventionally mean roll the other way) into a Pallas TPU kernel; parameter sweeps where shift is computed as a difference that can go negative.

Related errors

• stride and stride_axis must be both specified or not.
• 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/8731229d51f38c8d. Report an issue: GitHub.