{"record":{"id":"5a1183adde8d6ac2","repo":"jax-ml/jax","slug":"the-out-argument-to-jnp-stack-is-not-supported","errorCode":null,"errorMessage":"The 'out' argument to jnp.stack is not supported.","messagePattern":"The 'out' argument to jnp\\.stack is not supported\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/lax_numpy.py","lineNumber":4431,"sourceCode":"           [4, 5, 6]], dtype=int32)\n    >>> jnp.stack([x, y], axis=1)\n    Array([[1, 4],\n           [2, 5],\n           [3, 6]], dtype=int32)\n\n    :func:`~jax.numpy.unstack` performs the inverse operation:\n\n    >>> arr = jnp.stack([x, y], axis=1)\n    >>> x, y = jnp.unstack(arr, axis=1)\n    >>> x\n    Array([1, 2, 3], dtype=int32)\n    >>> y\n    Array([4, 5, 6], dtype=int32)\n  \"\"\"\n  if not len(arrays):\n    raise ValueError(\"Need at least one array to stack.\")\n  if out is not None:\n    raise NotImplementedError(\"The 'out' argument to jnp.stack is not supported.\")\n  if isinstance(arrays, (np.ndarray, Array)):\n    axis = _canonicalize_axis(axis, arrays.ndim)\n    return concatenate(expand_dims(arrays, axis + 1), axis=axis, dtype=dtype)\n  else:\n    arrays = util.ensure_arraylike_tuple(\"stack\", arrays)\n    if dtype is not None:\n      arrays = [asarray(a, dtype=dtype) for a in arrays]\n    else:\n      arrays = util.promote_dtypes(*arrays)\n    return lax.stack(arrays, axis=axis)\n\n\n@export\n@api.jit(static_argnames=\"axis\", inline=True)\ndef unstack(x: ArrayLike, /, *, axis: int = 0) -> tuple[Array, ...]:\n  \"\"\"Unstack an array along an axis.\n\n  JAX implementation of :func:`array_api.unstack`.","sourceCodeStart":4413,"sourceCodeEnd":4449,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/lax_numpy.py#L4413-L4449","documentation":"JAX arrays are immutable, so out= (in-place output) parameters from the numpy API are not supported; passing out to jnp.stack raises NotImplementedError.","triggerScenarios":"jnp.stack([a, b], axis=0, out=buf) — code ported verbatim from numpy that preallocates an output buffer.","commonSituations":"Porting numpy optimization patterns (preallocated buffers) to JAX; running numpy code under jax.","solutions":["Drop the out argument and assign the returned array","Use block-assignment style (x = x.at[...].set(...)) if you intended in-place writes"],"exampleFix":"// before\nout = jnp.empty((2, 3)); jnp.stack([a, b], out=out)\n// after\nout = jnp.stack([a, b])","handlingStrategy":"type-guard","validationCode":"assert out is None, 'out= is not supported in JAX'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Strip out= arguments when porting numpy code; use functional assignment"],"tags":["jnp-stack","out-argument","immutable-arrays","numpy-parity"],"backgroundTag":"unsupported-out-parameter","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}