jax-ml/jax · error · ValueError

The input must be non-scalar to take a cumulative sum, howev

Error message

The input must be non-scalar to take a cumulative sum, however a scalar value or scalar array was given.

What it means

jnp.cumulative_sum requires an array with at least one dimension; a scalar (0-d) input has no axis to accumulate over, so JAX raises ValueError immediately.

Source

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

  See Also:
    - :func:`jax.numpy.cumsum`: alternative API for cumulative sum.
    - :func:`jax.numpy.nancumsum`: cumulative sum while ignoring NaN values.
    - :func:`jax.numpy.add.accumulate`: cumulative sum via the ufunc API.

  Examples:
    >>> 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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Reshape to 1-d first: jnp.cumulative_sum(x.reshape(1)) or x[None]
  2. Check x.ndim > 0 before calling in generic code
  3. Use keepdims=True on the upstream reduction so the input stays non-scalar

Example fix

// before
jnp.cumulative_sum(jnp.sum(x))
// after
jnp.cumulative_sum(jnp.sum(x, keepdims=True))
Defensive patterns

Strategy: validation

Validate before calling

import jax.numpy as jnp

def cumulative_sum_safe(x, **kw):
    if jnp.ndim(x) == 0:
        x = jnp.reshape(x, (1,))
    return jnp.cumulative_sum(x, **kw)

Type guard

def is_nonscalar(x) -> bool:
    return jnp.asarray(x).ndim > 0

Prevention

When it happens

Trigger: Calling jnp.cumulative_sum(jnp.asarray(3.0)) or passing a Python scalar that becomes a 0-d array; x.ndim == 0.

Common situations: Feeding the output of an all-reduction (e.g. jnp.sum without keepdims) into cumulative_sum; looping over per-sample scalars in data pipelines.

Related errors


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