{"record":{"id":"8731229d51f38c8d","repo":"jax-ml/jax","slug":"shift-must-be-non-negative","errorCode":null,"errorMessage":"shift must be non-negative.","messagePattern":"shift must be non-negative\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/primitives.py","lineNumber":129,"sourceCode":"def _bitcast_batch_rule(batched_args, batch_axes, *, ty):\n  return bitcast(*batched_args, ty=ty), batch_axes[0]\n\nbatching.primitive_batchers[bitcast_p] = _bitcast_batch_rule\n\n\nroll_p = jax_core.Primitive(\"roll\")\n\n\ndef roll(\n    x: jax.Array,\n    shift: jax.Array | int,\n    axis: int,\n    *,\n    stride: int | None = None,\n    stride_axis: int | None = None,\n) -> jax.Array:\n  if isinstance(shift, int) and shift < 0:\n    raise ValueError(\"shift must be non-negative.\")\n  if axis < 0 or axis >= len(x.shape):\n    raise ValueError(\"axis is out of range.\")\n  if (stride is None) != (stride_axis is None):\n    raise ValueError(\"stride and stride_axis must be both specified or not.\")\n  if stride is not None and stride_axis is not None:\n    if stride < 0:\n      raise ValueError(\"stride must be non-negative.\")\n    if stride_axis < 0 or stride_axis >= len(x.shape):\n      raise ValueError(\"stride_axis is out of range\")\n    if axis == stride_axis:\n      raise ValueError(\"expected axis and stride_axis are different.\")\n  return roll_p.bind(\n      x, shift, axis=axis, stride=stride, stride_axis=stride_axis\n  )\n\n\n@roll_p.def_abstract_eval\ndef _roll_abstract_eval(x, shift, **_):","sourceCodeStart":111,"sourceCodeEnd":147,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/primitives.py#L111-L147","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Convert negative shift to its positive equivalent: shift % x.shape[axis]","Clamp or validate shift before calling roll","Compute shift = abs(shift) with direction handled by which end you read from, if applicable"],"exampleFix":"# before\ny = roll(x, shift=-3, axis=1)\n\n# after\ny = roll(x, shift=(-3) % x.shape[1], axis=1)","handlingStrategy":"validation","validationCode":"shift = shift % x.shape[axis] if isinstance(shift, int) else shift\nassert not (isinstance(shift, int) and shift < 0)\ny = roll(x, shift, axis=axis)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Normalize shifts with modulo the axis length","Never forward raw signed differences as shift"],"tags":["jax","pallas","mosaic","roll","argument-validation","shift"],"backgroundTag":"negative-shift-not-allowed","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}