jax-ml/jax · error · ValueError

XLA operations do not allow negative axes

Error message

XLA operations do not allow negative axes

What it means

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.

Source

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

def cumprod(operand: Array, axis: int = 0, reverse: bool = False) -> Array:
  """Computes a cumulative product along `axis`."""
  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):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Convert to a non-negative axis before calling: axis = axis % x.ndim
2. Use jnp.cumsum/jnp.nancumsum which accept negative axes
3. Upgrade JAX — newer versions normalize negative axes in lax.cum*

Example fix

// before
jax.lax.cumsum(x, axis=-1)
// after
jax.lax.cumsum(x, axis=x.ndim - 1)

axis = axis % x.ndim if axis < 0 else axis  # normalize before lax.cum*

When it happens

Trigger: Calling jax.lax.cumsum(x, axis=-1) (or cumprod/cummax/cummin/cumlogsumexp) with a negative axis on the raw lax primitive path.

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

Related errors

• axis {} is out of bounds for array of shape {}
• {} 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/80d2ed4430aa0a4d. Report an issue: GitHub.