jax-ml/jax · error · NotImplementedError

out argument of {self}

Error message

out argument of {self}

What it means

JAX ufunc objects (jnp.add, jnp.multiply, etc.) accept numpy's out keyword for signature compatibility but cannot write into a buffer because JAX arrays are immutable. Passing a non-None out raises NotImplementedError naming the ufunc.

Source

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

  def __hash__(self) -> int:
    # In both __hash__ and __eq__, we do not consider call, reduce, etc.
    # because they are considered implementation details rather than
    # necessary parts of object identity.
    return hash((self._func, self.__name__, self.identity,
                 self.nin, self.nout, self.nargs))

  def __eq__(self, other: Any) -> bool:
    return isinstance(other, ufunc) and (
      (self._func, self.__name__, self.identity, self.nin, self.nout, self.nargs) ==
      (other._func, other.__name__, other.identity, other.nin, other.nout, other.nargs))

  def __repr__(self) -> str:
    return f"<jnp.ufunc '{self.__name__}'>"

  def __call__(self, *args: ArrayLike, out: None = None, where: None = None) -> Any:
    check_arraylike(self.__name__, *args)
    if out is not None:
      raise NotImplementedError(f"out argument of {self}")
    if where is not None:
      raise NotImplementedError(f"where argument of {self}")
    call = (self.__static_props['call']
            or cast(Callable[..., Any], self._call_vectorized))
    return call(*args)

  @api.jit(static_argnames=['self'])
  def _call_vectorized(self, *args):
    return vectorize(self._func)(*args)

  @api.jit(static_argnames=['self', 'axis', 'dtype', 'out', 'keepdims'])
  def reduce(self, a: ArrayLike, axis: int | None = 0,
             dtype: DTypeLike | None = None,
             out: None = None, keepdims: bool = False, initial: ArrayLike | None = None,
             where: ArrayLike | None = None) -> Array:
    """Reduction operation derived from a binary function.

    JAX implementation of :meth:`numpy.ufunc.reduce`.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove the out argument and assign the result
2. Rewrite in-place accumulation patterns as functional updates (e.g. x = x + y or x.at[idx].set(...))

Example fix

// before
jnp.add(x, y, out=x)
// after
x = jnp.add(x, y)

kwargs.pop('out', None)  # before forwarding to a jnp.ufunc

When it happens

Trigger: jnp.add(x, y, out=buf) or any jnp.ufunc call with out=<array>.

Common situations: Ported numpy code using out= for in-place accumulation; generic wrapper code that forwards **kwargs including out.

Related errors

• The 'out' argument to jnp.outer is not supported.
• where argument of {self}
• out argument of {self.__name__}.reduce()
• out argument of {self.__name__}.accumulate()
• Because JAX arrays are immutable, jnp.ufunc.at() cannot oper

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