jax-ml/jax · error · ValueError

reduction axes {axes} contains out-of-bounds indices for {op

Error message

reduction axes {axes} contains out-of-bounds indices for {operand}.

What it means

A reduction primitive received an axis outside [0, operand.ndim). JAX checks every axis against the operand's rank before computing the output shape; both too-large positive and negative-beyond-rank indices fail.

Source

Thrown at jax/_src/lax/lax.py:8534

    raise TypeError("{} does not accept dtype {}. Accepted dtypes are subtypes "
                    "of number.".format(name, dtype_to_string(operand.dtype)))
  return operand.dtype

def _reduce_sum_transpose_rule(cotangent, operand, *, axes, out_sharding):
  assert ad.is_undefined_primal(operand)
  input_shape = operand.aval.shape
  broadcast_dimensions = tuple(np.delete(np.arange(len(input_shape)), axes))
  result = broadcast_in_dim(
      cotangent, input_shape, broadcast_dimensions,
      out_sharding=operand.aval.sharding)
  assert result.shape == input_shape
  return [result]

def _reduce_op_shape_rule(operand, *, axes, **_):
  if len(axes) != len(set(axes)):
    raise ValueError(f"duplicate value in 'axes' of reduction: {axes}")
  if not all(0 <= a < operand.ndim for a in axes):
    raise ValueError(f"reduction axes {axes} contains out-of-bounds indices for {operand}.")
  axes = frozenset(axes)
  return tuple(d for i, d in enumerate(operand.shape) if i not in axes)

def _reduce_op_sharding_rule_with_out_sharding(operand, *, axes, out_sharding):
  if out_sharding is not None:
    assert isinstance(out_sharding, NamedSharding)
    return out_sharding
  axes = frozenset(axes)
  new_spec = P(*tuple(s for i, s in enumerate(operand.sharding.spec.partitions)
                      if i not in axes))
  return operand.sharding.update(spec=new_spec)

def _reduce_op_unreduced_rule(operand, axes, out_sharding, out_kind, name):
  if out_sharding is not None and out_sharding.spec.unreduced:  # explicit mode
    if out_sharding.spec.unreduced_kind is not out_kind:
      raise core.ShardingTypeError(
          f"{name} requires `out_sharding`'s unreduced_kind to be {out_kind}"
          f' but got {out_sharding.spec}')

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Print x.ndim and the axes right before the call and clamp/validate: axes = tuple(a % x.ndim for a in axes).
2. If you removed a batch dim (squeeze/vmap), update hardcoded axis indices.
3. Validate user-supplied axis against the array rank: if not all(0 <= a < x.ndim for a in axes): raise ....
4. Use negative indexing intentionally (axis=-1 for last dim) instead of computing rank-dependent positives.

Example fix

# before
out = lax.reduce_sum(x, axes=(2,))  # x is 2-D
# after
out = lax.reduce_sum(x, axes=(-1,))  # last axis, rank-independent

assert all(0 <= a < x.ndim or -x.ndim <= a < 0 for a in axes), (x.shape, axes)
axes = tuple(a % x.ndim for a in axes)
out = lax.reduce_sum(x, axes)

def axes_in_bounds(x, axes):
    return all(-x.ndim <= a < x.ndim for a in axes)

When it happens

Trigger: lax.reduce_max(x, axes=(2,)) on a 2-D array; jnp.sum(x, axis=3) where x.ndim==2; passing axis=-3 to a 2-D array via jnp.prod. Common when axis is computed from user input or config.

Common situations: Hardcoded axes that stop matching after adding/removing a batch dimension; loops that assume a fixed rank; passing a numpy-style axis that exceeds rank after jnp.squeeze/vmap transformations change the number of dimensions.

Related errors

• duplicate value in 'axes' of reduction: {axes}
• axis argument out of range: {axis=} for {operand.shape=}
• 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/5825cf2fbc643f01. Report an issue: GitHub.