jax-ml/jax · error · NotImplementedError

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

Error message

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

What it means

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.

Source

Thrown at jax/_src/numpy/indexing.py:704

    example, we can instead clip to the last valid value:

    >>> jnp.take(x, indices, axis=0, mode='clip')
    Array([[4., 5., 6.],
           [1., 2., 3.]], dtype=float32)
    >>> x.at[indices].get(mode='clip')  # equivalent indexing syntax
    Array([[4., 5., 6.],
           [1., 2., 3.]], dtype=float32)
  """
  return _take(a, indices, None if axis is None else operator.index(axis), out,
               mode, unique_indices=unique_indices, indices_are_sorted=indices_are_sorted,
               fill_value=fill_value)


@api.jit(static_argnames=('axis', 'mode', 'unique_indices', 'indices_are_sorted', 'fill_value'))
def _take(a, indices, axis: int | None = None, out=None, mode=None,
          unique_indices=False, indices_are_sorted=False, fill_value=None):
  if out is not None:
    raise NotImplementedError("The 'out' argument to jnp.take is not supported.")
  a, indices = util.ensure_arraylike("take", a, indices)

  if axis is None:
    a = a.ravel()
    axis_idx = 0
  else:
    axis_idx = canonicalize_axis(axis, np.ndim(a))

  if mode is None or mode == "fill":
    gather_mode = slicing.GatherScatterMode.FILL_OR_DROP
    # lax.gather() does not support negative indices, so we wrap them here
    indices = util._where(indices < 0, indices + a.shape[axis_idx], indices)
  elif mode == "raise":
    # TODO(phawkins): we have no way to report out of bounds errors yet.
    raise NotImplementedError("The 'raise' mode to jnp.take is not supported.")
  elif mode == "wrap":
    indices = ufuncs.mod(indices, lax._const(indices, a.shape[axis_idx]))
    gather_mode = slicing.GatherScatterMode.PROMISE_IN_BOUNDS

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove the out= argument and use the returned array: y = jnp.take(a, indices)
2. If buffer reuse matters, manage buffers outside JAX (e.g. with donation in jax.jit)

Example fix

// before
out = np.empty(n)
jnp.take(a, idx, out=out)
// after
out = jnp.take(a, idx)

assert out is None, 'jnp.take does not support out='

When it happens

Trigger: Calling jnp.take(a, indices, out=buf) with any non-None out argument.

Common situations: Porting NumPy code that reuses a preallocated output buffer via out=; performance-oriented NumPy idioms copied into JAX.

Related errors

• array() takes at most 5 positional arguments but {len(args)
• array() got multiple values for argument '{name}'
• Only implemented for order='K'
• np.reshape order=A is not implemented.
• `type` argument of array.view() is not supported.

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