jax-ml/jax · error · NotImplementedError

The 'out' argument to jnp.stack is not supported.

Error message

The 'out' argument to jnp.stack is not supported.

What it means

JAX arrays are immutable, so out= (in-place output) parameters from the numpy API are not supported; passing out to jnp.stack raises NotImplementedError.

Source

Thrown at jax/_src/numpy/lax_numpy.py:4431

           [4, 5, 6]], dtype=int32)
    >>> jnp.stack([x, y], axis=1)
    Array([[1, 4],
           [2, 5],
           [3, 6]], dtype=int32)

    :func:`~jax.numpy.unstack` performs the inverse operation:

    >>> arr = jnp.stack([x, y], axis=1)
    >>> x, y = jnp.unstack(arr, axis=1)
    >>> x
    Array([1, 2, 3], dtype=int32)
    >>> y
    Array([4, 5, 6], dtype=int32)
  """
  if not len(arrays):
    raise ValueError("Need at least one array to stack.")
  if out is not None:
    raise NotImplementedError("The 'out' argument to jnp.stack is not supported.")
  if isinstance(arrays, (np.ndarray, Array)):
    axis = _canonicalize_axis(axis, arrays.ndim)
    return concatenate(expand_dims(arrays, axis + 1), axis=axis, dtype=dtype)
  else:
    arrays = util.ensure_arraylike_tuple("stack", arrays)
    if dtype is not None:
      arrays = [asarray(a, dtype=dtype) for a in arrays]
    else:
      arrays = util.promote_dtypes(*arrays)
    return lax.stack(arrays, axis=axis)


@export
@api.jit(static_argnames="axis", inline=True)
def unstack(x: ArrayLike, /, *, axis: int = 0) -> tuple[Array, ...]:
  """Unstack an array along an axis.

  JAX implementation of :func:`array_api.unstack`.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Drop the out argument and assign the returned array
2. Use block-assignment style (x = x.at[...].set(...)) if you intended in-place writes

Example fix

// before
out = jnp.empty((2, 3)); jnp.stack([a, b], out=out)
// after
out = jnp.stack([a, b])

assert out is None, 'out= is not supported in JAX'

When it happens

Trigger: jnp.stack([a, b], axis=0, out=buf) — code ported verbatim from numpy that preallocates an output buffer.

Common situations: Porting numpy optimization patterns (preallocated buffers) to JAX; running numpy code under jax.

Related errors

• The 'out' argument to jnp.choose is not supported.
• The 'out' argument to jnp.std is not supported.
• Value of type {type(self)} is not indexable.
• Need at least one array to stack.
• The 'out' argument to jnp.{name} is not supported.

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