{"record":{"id":"80d2ed4430aa0a4d","repo":"jax-ml/jax","slug":"xla-operations-do-not-allow-negative-axes","errorCode":null,"errorMessage":"XLA operations do not allow negative axes","messagePattern":"XLA operations do not allow negative axes","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/control_flow/loops.py","lineNumber":3055,"sourceCode":"def cumprod(operand: Array, axis: int = 0, reverse: bool = False) -> Array:\n  \"\"\"Computes a cumulative product along `axis`.\"\"\"\n  return cumprod_p.bind(operand, axis=int(axis), reverse=bool(reverse))\n\ndef cummax(operand: Array, axis: int = 0, reverse: bool = False) -> Array:\n  \"\"\"Computes a cumulative maximum along `axis`.\"\"\"\n  return cummax_p.bind(operand, axis=int(axis), reverse=bool(reverse))\n\ndef cummin(operand: Array, axis: int = 0, reverse: bool = False) -> Array:\n  \"\"\"Computes a cumulative minimum along `axis`.\"\"\"\n  return cummin_p.bind(operand, axis=int(axis), reverse=bool(reverse))\n\ndef cumlogsumexp(operand: Array, axis: int = 0, reverse: bool = False) -> Array:\n  \"\"\"Computes a cumulative logsumexp along `axis`.\"\"\"\n  return cumlogsumexp_p.bind(operand, axis=int(axis), reverse=bool(reverse))\n\ndef _cumred_shape_rule(x, *, axis: int, reverse: bool):\n  if axis < 0:\n    raise ValueError(\"XLA operations do not allow negative axes\")\n  elif axis >= x.ndim:\n    raise ValueError(\n        f\"axis {axis} is out of bounds for array of shape {x.shape}\")\n  return x.shape\n\ndef _cumred_sharding_rule(x, *, axis: int, reverse: bool):\n  if x.sharding.spec[axis] is not None:\n    raise core.ShardingTypeError(\n        'Input should be unsharded over the axis being reduced. Got input'\n        f' type={x} and {axis=}')\n  return x.sharding\n\ndef _cumsum_transpose_rule(t, operand, *, axis: int, reverse: bool):\n  return [cumsum(t, axis=axis, reverse=not reverse)]\n\n\ndef cumred_reduce_window_impl(window_reduce: Callable, x, *, axis: int,\n                              reverse: bool):","sourceCodeStart":3037,"sourceCodeEnd":3073,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/control_flow/loops.py#L3037-L3073","documentation":"The cumulative-reduction primitives (cumsum, cumprod, cumlogsumexp, cummax, cummin) bind into XLA ReduceWindow ops, which (per this rule) require non-negative axis indices. The shape rule rejects negative axes even though Python/numpy conventionally allows them, so axis=-1 raises ValueError.","triggerScenarios":"Calling jax.lax.cumsum(x, axis=-1) (or cumprod/cummax/cummin/cumlogsumexp) with a negative axis on the raw lax primitive path.","commonSituations":"Copy-pasting numpy-style code that uses axis=-1; note jnp.cumsum normalizes negative axes, so this typically appears only via jax.lax.cum* with a negative axis, or versions where normalization differs.","solutions":["Convert to a non-negative axis before calling: axis = axis % x.ndim","Use jnp.cumsum/jnp.nancumsum which accept negative axes","Upgrade JAX — newer versions normalize negative axes in lax.cum*"],"exampleFix":"// before\njax.lax.cumsum(x, axis=-1)\n// after\njax.lax.cumsum(x, axis=x.ndim - 1)","handlingStrategy":"validation","validationCode":"axis = axis % x.ndim if axis < 0 else axis  # normalize before lax.cum*","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Prefer jnp.cumsum which normalizes negative axes","Wrap lax.cum* calls with a small normalize_axis helper"],"tags":["jax","lax","cumsum","negative-axis"],"backgroundTag":"negative-axis-index","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}