jax-ml/jax · error · ValueError

The input array has rank {x.ndim}, however axis was not set

Error message

The input array has rank {x.ndim}, however axis was not set to an explicit value. The axis argument is only optional for one-dimensional arrays.

What it means

For jnp.cumulative_sum, the axis argument is only optional for 1-d arrays. If the input has rank > 1 and axis is None, JAX raises ValueError telling you to specify the axis explicitly (unlike cumsum, it will not flatten).

Source

Thrown at jax/_src/numpy/reductions.py:2308

    >>> x = jnp.array([[1, 2, 3],
    ...                [4, 5, 6]])
    >>> jnp.cumulative_sum(x, axis=1)
    Array([[ 1,  3,  6],
           [ 4,  9, 15]], dtype=int32)
    >>> jnp.cumulative_sum(x, axis=1, include_initial=True)
    Array([[ 0,  1,  3,  6],
           [ 0,  4,  9, 15]], dtype=int32)
  """
  x = ensure_arraylike("cumulative_sum", x)
  if x.ndim == 0:
    raise ValueError(
      "The input must be non-scalar to take a cumulative sum, however a "
      "scalar value or scalar array was given."
    )
  if axis is None:
    axis = 0
    if x.ndim > 1:
      raise ValueError(
        f"The input array has rank {x.ndim}, however axis was not set to an "
        "explicit value. The axis argument is only optional for one-dimensional "
        "arrays.")

  axis = canonicalize_axis(axis, x.ndim)
  if dtype is not None:
    dtype = dtypes.check_and_canonicalize_user_dtype(dtype)
  out = _cumsum_with_promotion(x, axis=axis, dtype=dtype)
  if include_initial:
    zeros_shape = list(x.shape)
    zeros_shape[axis] = 1
    out = lax.concatenate(
      [lax.full(zeros_shape, 0, dtype=out.dtype), out],
      dimension=axis)
  return out


@export

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Pass an explicit axis: jnp.cumulative_sum(x, axis=1)
Use jnp.cumsum if flatten-over-all-elements semantics are actually wanted

Example fix

// before
jnp.cumulative_sum(batch)  # batch.ndim == 2
// after
jnp.cumulative_sum(batch, axis=1)

Defensive patterns

Strategy: validation

Validate before calling

import jax.numpy as jnp

def cumulative_sum_safe(x, axis=None, **kw):
    if axis is None and jnp.ndim(x) > 1:
        raise ValueError('specify axis for multi-dim cumulative_sum')
    return jnp.cumulative_sum(x, axis=axis, **kw)

Prevention

Always pass axis for batched inputs
Write tests with 2-d fixtures, not just 1-d
Remember cumulative_sum does not flatten like cumsum

When it happens

Trigger: Calling jnp.cumulative_sum(x) where x.ndim > 1, e.g. a (3, 4) matrix with axis omitted.

Common situations: Assuming cumulative_sum behaves like np.cumsum (which flattens by default) or like cumsum on the last axis; testing with 1-d data then deploying on batches.

Related errors

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