{"record":{"id":"21c145cd69fcf1ad","repo":"jax-ml/jax","slug":"axis-is-out-of-range","errorCode":null,"errorMessage":"axis is out of range.","messagePattern":"axis is out of range\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/primitives.py","lineNumber":131,"sourceCode":"\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, **_):\n  del shift\n  return x","sourceCodeStart":113,"sourceCodeEnd":149,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/primitives.py#L113-L149","documentation":"roll in Mosaic validates that the roll axis is a valid non-negative index into x's shape (negative indices are not normalized like jnp.roll). Passing axis < 0 or axis >= x.ndim raises this ValueError immediately.","triggerScenarios":"Calling mosaic roll with a negative axis (roll(x, 2, axis=-1)) or an axis index equal to or beyond the rank of x.","commonSituations":"Translating jnp.roll or np.roll calls that idiomatically use axis=-1/-2 into a Pallas kernel; changing tensor rank (e.g. adding a batch dim) so a previously valid axis constant is now out of range.","solutions":["Normalize negative axes first: axis = axis % len(x.shape)","Assert axis is in range before calling roll","Use the last-axis spelling explicitly, e.g. axis=x.ndim - 1 instead of -1"],"exampleFix":"# before\ny = roll(x, 2, axis=-1)\n\n# after\ny = roll(x, 2, axis=len(x.shape) - 1)","handlingStrategy":"validation","validationCode":"assert 0 <= axis < len(x.shape), f\"axis {axis} out of range for shape {x.shape}\"\n# normalize negative axes like jnp:\naxis = axis % len(x.shape)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Normalize negative axes before Pallas calls","Re-verify axis constants whenever tensor rank changes"],"tags":["jax","pallas","mosaic","roll","axis-validation","shape-validation"],"backgroundTag":"axis-out-of-range","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}