jax-ml/jax · error · ValueError

axis {} is out of bounds for array of shape {}

Error message

axis {} is out of bounds for array of shape {}

What it means

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.

Source

Thrown at jax/_src/lax/control_flow/loops.py:3057

  return cumprod_p.bind(operand, axis=int(axis), reverse=bool(reverse))

def cummax(operand: Array, axis: int = 0, reverse: bool = False) -> Array:
  """Computes a cumulative maximum along `axis`."""
  return cummax_p.bind(operand, axis=int(axis), reverse=bool(reverse))

def cummin(operand: Array, axis: int = 0, reverse: bool = False) -> Array:
  """Computes a cumulative minimum along `axis`."""
  return cummin_p.bind(operand, axis=int(axis), reverse=bool(reverse))

def cumlogsumexp(operand: Array, axis: int = 0, reverse: bool = False) -> Array:
  """Computes a cumulative logsumexp along `axis`."""
  return cumlogsumexp_p.bind(operand, axis=int(axis), reverse=bool(reverse))

def _cumred_shape_rule(x, *, axis: int, reverse: bool):
  if axis < 0:
    raise ValueError("XLA operations do not allow negative axes")
  elif axis >= x.ndim:
    raise ValueError(
        f"axis {axis} is out of bounds for array of shape {x.shape}")
  return x.shape

def _cumred_sharding_rule(x, *, axis: int, reverse: bool):
  if x.sharding.spec[axis] is not None:
    raise core.ShardingTypeError(
        'Input should be unsharded over the axis being reduced. Got input'
        f' type={x} and {axis=}')
  return x.sharding

def _cumsum_transpose_rule(t, operand, *, axis: int, reverse: bool):
  return [cumsum(t, axis=axis, reverse=not reverse)]


def cumred_reduce_window_impl(window_reduce: Callable, x, *, axis: int,
                              reverse: bool):
  n = x.shape[axis]
  if n == 0:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Check x.ndim and clamp: axis = min(axis, x.ndim - 1)
2. Reshape the operand so the intended axis exists: x[:, None] or x[None, :]
3. Validate the input rank upstream (assert x.ndim == expected)

Example fix

// before
jax.lax.cumsum(x, axis=1)  # x is shape (n,)
// after
jax.lax.cumsum(x.reshape(n, 1), axis=0)  # or fix data pipeline to keep rank 2

assert -x.ndim <= axis < x.ndim, f'axis {axis} invalid for rank {x.ndim}'

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• XLA operations do not allow negative axes
• {} does not accept dtype {}. Accepted dtypes are subtypes of
• lax.platform_dependent: the '{pname}' branch must be a calla
• Use 'cuda', 'rocm', or 'oneapi' for lax.platform_dependent.
• lax.platform_dependent: the 'default' branch must be a calla

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/c970f4596299ee83. Report an issue: GitHub.