jax-ml/jax · error · ValueError

accumulate does not allow multiple axes

Error message

accumulate does not allow multiple axes

What it means

ufunc.accumulate operates along exactly one axis; numpy's semantics of axis=None (flatten) or a tuple of axes are rejected by JAX's scan-based accumulate with this ValueError.

Source

Thrown at jax/_src/numpy/ufunc_api.py:390

      raise ValueError("accumulate only supported for binary ufuncs")
    if self.nout != 1:
      raise ValueError("accumulate only supported for functions returning a single value")
    if out is not None:
      raise NotImplementedError(f"out argument of {self.__name__}.accumulate()")
    accumulate = self.__static_props['accumulate'] or self._accumulate_via_scan
    return accumulate(a, axis=axis, dtype=dtype)

  def _accumulate_via_scan(self, arr: ArrayLike, axis: int = 0,
                           dtype: DTypeLike | None = None) -> Array:
    assert self.nin == 2 and self.nout == 1
    check_arraylike(f"{self.__name__}.accumulate", arr)
    arr = lax.asarray(arr)

    if dtype is None:
      dtype = api.eval_shape(self._func, lax._one(arr), lax._one(arr)).dtype

    if axis is None or isinstance(axis, tuple):
      raise ValueError("accumulate does not allow multiple axes")
    axis = canonicalize_axis(axis, np.ndim(arr))

    if arr.size == 0:
      return lax.full(arr.shape, 0, dtype)
    arr = _moveaxis(arr, axis, 0)
    def scan_fun(carry, _):
      i, x = carry
      y = _where(i == 0, arr[0].astype(dtype), self(x.astype(dtype), arr[i].astype(dtype)))
      return (i + 1, y), y
    _, result = control_flow.scan(scan_fun, (0, arr[0].astype(dtype)), None, length=arr.shape[0])
    return _moveaxis(result, 0, axis)

  @api.jit(static_argnums=[0], static_argnames=['inplace'])
  def at(self, a: ArrayLike, indices: Any, b: ArrayLike | None = None, /, *,
         inplace: bool = True) -> Array:
    """Update elements of an array via the specified unary or binary ufunc.

    JAX implementation of :func:`numpy.ufunc.at`.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass a single integer axis
2. For axis=None behavior, ravel first: jnp.add.accumulate(x.ravel()) then reshape back

Example fix

// before
jnp.add.accumulate(x, axis=None)
// after
jnp.add.accumulate(x.ravel()).reshape(x.shape)

assert isinstance(axis, int), 'accumulate requires a single int axis'

When it happens

Trigger: jnp.add.accumulate(x, axis=None) or jnp.add.accumulate(x, axis=(0,1)).

Common situations: Porting numpy code using axis=None to accumulate over a flattened array; generic code that forwards tuple axes.

Related errors

• accumulate only supported for binary ufuncs
• accumulate only supported for functions returning a single v
• out argument of {self.__name__}.accumulate()
• reduction axes {axes} contains out-of-bounds indices for {op
• Invalid axis {axis} for operand shape {operand.shape}

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