jax-ml/jax · error · NotImplementedError

The 'out' argument to jnp.{name} is not supported

Error message

The 'out' argument to jnp.{name} is not supported

What it means

The shared helper behind jnp.cumsum, cumprod, nancumsum, and nancumprod rejects the out= argument because JAX arrays are immutable and cannot be filled in place.

Source

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

  a = ensure_arraylike("nanstd", a)
  where = check_where("nanstd", where)
  if dtype is not None:
    dtype = dtypes.check_and_canonicalize_user_dtype(dtype, "nanstd")
  if out is not None:
    raise NotImplementedError("The 'out' argument to jnp.nanstd is not supported.")
  return lax.sqrt(nanvar(a, axis=axis, dtype=dtype, ddof=ddof,
                         keepdims=keepdims, where=where, mean=mean))


def _cumulative_reduction(
    name: str, reduction: Callable[..., Array],
    a: ArrayLike, axis: int | None, dtype: DTypeLike | None, out: None = None,
    fill_nan: bool = False, fill_value: ArrayLike = 0,
    promote_integers: bool = False) -> Array:
  """Helper function for implementing cumulative reductions."""
  a = ensure_arraylike(name, a)
  if out is not None:
    raise NotImplementedError(f"The 'out' argument to jnp.{name} is not supported")

  if axis is None or _isscalar(a):
    if not builtins.all(s is None for s in core.typeof(a).sharding.spec):
      raise core.ShardingTypeError(
          "The input should be fully replicated when axis is not specified to"
          f" {name}. Got input type={core.typeof(a)}")
    a = lax.reshape(a, (np.size(a),))
  if axis is None:
    axis = 0

  a_shape = list(np.shape(a))
  num_dims = len(a_shape)
  axis = canonicalize_axis(axis, num_dims)

  if fill_nan:
    a = _where(lax._isnan(a), lax._const(a, fill_value), a)

  computation_type: DType

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Drop out= and use the return value
2. Preallocate nothing — rely on jit to fuse and manage memory

Example fix

// before
np.cumsum(x, out=buf)
// after
buf = jnp.cumsum(x)

def cumulate(fn, x, **kw):
    kw.pop('out', None)
    return fn(x, **kw)

When it happens

Trigger: Calling jnp.cumsum(x, out=buf), jnp.cumprod(x, out=buf), jnp.nancumsum(x, out=buf), or jnp.nancumprod(x, out=buf).

Common situations: Porting NumPy cumulative-reduction code that preallocated output buffers; generic **kwargs forwarding from array-API shims.

Related errors

• The 'out' argument to jnp.std is not supported.
• The 'out' argument to jnp.ptp is not supported.
• The 'out' argument to jnp.nanmean is not supported.
• The 'out' argument to jnp.nanvar is not supported.
• The 'out' argument to jnp.nanstd is not supported.

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