jax-ml/jax · error · NotImplementedError

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

Error message

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

What it means

JAX arrays are immutable, so argmin cannot deposit its result index into a caller-supplied out array. The wrapper rejects any non-None out argument to jnp.argmin.

Solutions

  1. Drop out and assign the return value: idx = jnp.argmin(a)
  2. Simulate buffer reuse via buf.at[:].set(jnp.argmin(a)) if truly needed
  3. Filter out= from kwargs before delegating to jnp

Example fix

// before
jnp.argmin(a, axis=0, out=buf)
// after
buf = jnp.argmin(a, axis=0)
Defensive patterns

Strategy: validation

Validate before calling

kwargs.pop('out', None)  # before calling jnp.argmin

Prevention

When it happens

Trigger: Calling jnp.argmin(a, out=buf), typically from ported NumPy code that preallocated an output buffer.

Common situations: NumPy performance idioms using out=; generic **kwargs forwarding wrappers that pass out through.

Related errors

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

Appendix: source

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

    - :func:`jax.numpy.nanargmin`: compute ``argmin`` while ignoring NaN values.

  Examples:
    >>> x = jnp.array([1, 3, 5, 4, 2])
    >>> jnp.argmin(x)
    Array(0, dtype=int32)

    >>> x = jnp.array([[1, 3, 2],
    ...                [5, 4, 1]])
    >>> jnp.argmin(x, axis=1)
    Array([0, 2], dtype=int32)

    >>> jnp.argmin(x, axis=1, keepdims=True)
    Array([[0],
           [2]], dtype=int32)
  """
  arr = util.ensure_arraylike("argmin", a)
  if out is not None:
    raise NotImplementedError("The 'out' argument to jnp.argmin is not supported.")
  return _argmin(arr, None if axis is None else operator.index(axis),
                 keepdims=bool(keepdims))

@api.jit(static_argnames=('axis', 'keepdims'), inline=True)
def _argmin(a: Array, axis: int | None = None, keepdims: bool = False) -> Array:
  if axis is None:
    dims = list(range(np.ndim(a)))
    a = ravel(a)
    axis = 0
  else:
    dims = [axis]
  if a.shape[axis] == 0:
    raise ValueError("attempt to get argmin of an empty sequence")
  # TODO(phawkins): use an int64 index if the dimension is large enough.
  result = lax.argmin(a, _canonicalize_axis(axis, a.ndim), int)
  return expand_dims(result, dims) if keepdims else result

View on GitHub (pinned to 1e1c6a8fc0)