{"record":{"id":"c970f4596299ee83","repo":"jax-ml/jax","slug":"axis-is-out-of-bounds-for-array-of-shape","errorCode":null,"errorMessage":"axis {} is out of bounds for array of shape {}","messagePattern":"axis (.+?) is out of bounds for array of shape (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/control_flow/loops.py","lineNumber":3057,"sourceCode":"  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):\n  n = x.shape[axis]\n  if n == 0:","sourceCodeStart":3039,"sourceCodeEnd":3075,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/control_flow/loops.py#L3039-L3075","documentation":"The cumulative-reduction shape rule validates that the axis falls within the operand's rank; an axis >= x.ndim cannot name any dimension so XLA could not lower the op, and a ValueError with the offending axis and shape is raised.","triggerScenarios":"Calling jax.lax.cumsum(x, axis=1) on a 1-D array, or generally axis >= x.ndim; often after a squeeze/jnp.ravel removed the intended dimension.","commonSituations":"Applying cumsum over a batch axis that was accidentally squeezed; assuming input is 2-D (e.g. (batch, time)) when it is actually 1-D (time,) after tree mapping over leaves.","solutions":["Check x.ndim and clamp: axis = min(axis, x.ndim - 1)","Reshape the operand so the intended axis exists: x[:, None] or x[None, :]","Validate the input rank upstream (assert x.ndim == expected)"],"exampleFix":"// before\njax.lax.cumsum(x, axis=1)  # x is shape (n,)\n// after\njax.lax.cumsum(x.reshape(n, 1), axis=0)  # or fix data pipeline to keep rank 2","handlingStrategy":"validation","validationCode":"assert -x.ndim <= axis < x.ndim, f'axis {axis} invalid for rank {x.ndim}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Check x.ndim after squeeze/ravel operations","Clamp axis: axis = min(axis, x.ndim - 1)"],"tags":["jax","lax","cumsum","axis-out-of-bounds"],"backgroundTag":"axis-out-of-bounds","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}