{"record":{"id":"25ba4f1ea14f5041","repo":"jax-ml/jax","slug":"the-out-argument-to-jnp-take-is-not-supported","errorCode":null,"errorMessage":"The 'out' argument to jnp.take is not supported.","messagePattern":"The 'out' argument to jnp\\.take is not supported\\.","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/indexing.py","lineNumber":704,"sourceCode":"    example, we can instead clip to the last valid value:\n\n    >>> jnp.take(x, indices, axis=0, mode='clip')\n    Array([[4., 5., 6.],\n           [1., 2., 3.]], dtype=float32)\n    >>> x.at[indices].get(mode='clip')  # equivalent indexing syntax\n    Array([[4., 5., 6.],\n           [1., 2., 3.]], dtype=float32)\n  \"\"\"\n  return _take(a, indices, None if axis is None else operator.index(axis), out,\n               mode, unique_indices=unique_indices, indices_are_sorted=indices_are_sorted,\n               fill_value=fill_value)\n\n\n@api.jit(static_argnames=('axis', 'mode', 'unique_indices', 'indices_are_sorted', 'fill_value'))\ndef _take(a, indices, axis: int | None = None, out=None, mode=None,\n          unique_indices=False, indices_are_sorted=False, fill_value=None):\n  if out is not None:\n    raise NotImplementedError(\"The 'out' argument to jnp.take is not supported.\")\n  a, indices = util.ensure_arraylike(\"take\", a, indices)\n\n  if axis is None:\n    a = a.ravel()\n    axis_idx = 0\n  else:\n    axis_idx = canonicalize_axis(axis, np.ndim(a))\n\n  if mode is None or mode == \"fill\":\n    gather_mode = slicing.GatherScatterMode.FILL_OR_DROP\n    # lax.gather() does not support negative indices, so we wrap them here\n    indices = util._where(indices < 0, indices + a.shape[axis_idx], indices)\n  elif mode == \"raise\":\n    # TODO(phawkins): we have no way to report out of bounds errors yet.\n    raise NotImplementedError(\"The 'raise' mode to jnp.take is not supported.\")\n  elif mode == \"wrap\":\n    indices = ufuncs.mod(indices, lax._const(indices, a.shape[axis_idx]))\n    gather_mode = slicing.GatherScatterMode.PROMISE_IN_BOUNDS","sourceCodeStart":686,"sourceCodeEnd":722,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/indexing.py#L686-L722","documentation":"jnp.take does not support the out= parameter because JAX arrays are immutable and out-of-place updates cannot write into a caller-supplied buffer under jit transformations.","triggerScenarios":"Calling jnp.take(a, indices, out=buf) with any non-None out argument.","commonSituations":"Porting NumPy code that reuses a preallocated output buffer via out=; performance-oriented NumPy idioms copied into JAX.","solutions":["Remove the out= argument and use the returned array: y = jnp.take(a, indices)","If buffer reuse matters, manage buffers outside JAX (e.g. with donation in jax.jit)"],"exampleFix":"// before\nout = np.empty(n)\njnp.take(a, idx, out=out)\n// after\nout = jnp.take(a, idx)","handlingStrategy":"validation","validationCode":"assert out is None, 'jnp.take does not support out='","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Never pass out= to jnp functions; JAX arrays are immutable","Use returned values and jit buffer donation for reuse"],"tags":["jax","numpy-compat","out-arg","immutable"],"backgroundTag":"unsupported-out-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}