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

jnp.argmin does not support the out parameter; JAX arrays are immutable and functional, so writing results into a caller-supplied buffer is unsupported, unlike NumPy.

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)

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)

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

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

• The 'out' argument to jnp.argmax is not supported.
• The 'out' argument to jnp.nanargmax is not supported.
• The 'out' argument to jnp.nanargmin is not supported.
• The 'out' argument to jnp.compress is not supported.
• 'order' argument to argsort is not supported.

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