{"record":{"id":"0e781dc3698e5e4f","repo":"jax-ml/jax","slug":"expected-axis-and-stride-axis-are-different","errorCode":null,"errorMessage":"expected axis and stride_axis are different.","messagePattern":"expected axis and stride_axis are different\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/primitives.py","lineNumber":140,"sourceCode":"    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\n\n\ndef _roll_lowering_rule(\n    ctx: mlir.LoweringRuleContext, x, shift, *, axis, stride, stride_axis\n):\n  def _roll(x, shift):\n    if stride is None:\n      return jnp.roll(x, shift, axis)\n    outputs = [","sourceCodeStart":122,"sourceCodeEnd":158,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/primitives.py#L122-L158","documentation":"Strided roll rolls along `axis` while stepping along `stride_axis`; these must be different dimensions for the operation to be meaningful. Passing the same axis for both is rejected because the stride would fold into the roll itself.","triggerScenarios":"roll(x, shift, axis=1, stride=4, stride_axis=1).","commonSituations":"Defaulting both parameters to the same value in a wrapper; refactoring where stride_axis was copied from axis and never changed; misunderstanding the two-axis strided rolling API.","solutions":["Choose a stride_axis different from axis (typically the dimension you are iterating blocks over)","If you only want a plain roll, drop stride and stride_axis entirely","Add an assert axis != stride_axis in your kernel wrapper for early failure"],"exampleFix":"# before\ny = roll(x, 2, axis=1, stride=4, stride_axis=1)\n\n# after\ny = roll(x, 2, axis=1, stride=4, stride_axis=0)\n# or plain roll:\ny = roll(x, 2, axis=1)","handlingStrategy":"validation","validationCode":"if stride is not None:\n  assert axis != stride_axis, \"axis and stride_axis must differ\"\n# drop strided args entirely when unused:\nif stride is None:\n  y = roll(x, shift, axis=axis)\nelse:\n  y = roll(x, shift, axis=axis, stride=stride, stride_axis=stride_axis)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Default stride_axis to a different dimension than axis","Add the inequality assert in wrapper functions"],"tags":["jax","pallas","mosaic","roll","argument-validation"],"backgroundTag":"duplicate-axis-arguments","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}