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

JAX does not support the numpy 'out' parameter for reductions because JAX arrays are immutable and outputs cannot be written in place. The parameter is accepted only as None for numpy-delegation compatibility.

Source

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

ReductionOp = Callable[[Any, Any], Any]

def _reduction(a: ArrayLike, name: str, op: ReductionOp, init_val: ArrayLike,
               *, has_identity: bool = True,
               preproc: Callable[[Array], Array] | None = None,
               bool_op: ReductionOp | None = None,
               upcast_f16_for_computation: bool = False,
               axis: Axis = None, dtype: DTypeLike | None = None, out: None = None,
               keepdims: bool = False, initial: ArrayLike | None = None,
               where_: ArrayLike | None = None,
               parallel_reduce: Callable[..., Array] | None = None,
               promote_integers: bool = False) -> Array:
  bool_op = bool_op or op
  # Note: we must accept out=None as an argument, because numpy reductions delegate to
  # object methods. For example `np.sum(x)` will call `x.sum()` if the `sum()` method
  # exists, passing along all its arguments.
  if out is not None:
    raise NotImplementedError(f"The 'out' argument to jnp.{name} is not supported.")
  a = ensure_arraylike(name, a)
  where_ = check_where(name, where_)
  axis = core.concrete_or_error(None, axis, f"axis argument to jnp.{name}().")

  if initial is None and not has_identity and where_ is not None:
    raise ValueError(f"reduction operation {name} does not have an identity, so to use a "
                     f"where mask one has to specify 'initial'")

  a = preproc(a) if preproc else a
  pos_dims, dims = _reduction_dims(a, axis)

  if initial is None and not has_identity:
    shape = np.shape(a)
    if not _all(shape[d] >= 1 for d in pos_dims):
      raise ValueError(f"zero-size array to reduction operation {name} which has no identity")

  result_dtype: DType
  if dtype is None:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Drop the out argument and use the returned value: y = jnp.sum(x)
2. Use y = y.at[...].set(jnp.sum(x)) if you need to update a buffer under jit
3. Replace numpy namespace with jax.numpy consistently so out is never threaded through

Example fix

// before
jnp.sum(x, out=result)
// after
result = jnp.sum(x)

kwargs = {k: v for k, v in kwargs.items() if k != 'out'}
jnp.sum(x, **kwargs)

When it happens

Trigger: Calling jnp.sum(x, out=buf), np.sum(jax_array, out=...) (numpy delegates to the array's .sum method), or any reduction (prod/max/min/all/any) with a non-None out argument.

Common situations: Porting numpy code that reuses output buffers for performance; code using np.add.reduce(x, out=y) on JAX arrays via __array_ufunc__ delegation.

Related errors

• Value of type {type(self)} is not indexable.
• The 'out' argument to jnp.round is not supported.
• The 'out' argument to jnp.compress is not supported.
• jnp.{name}: where must be None or a boolean array; got {wher
• reduction operation {name} does not have an identity, so to

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