jax-ml/jax · error · NotImplementedError

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

Error message

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

What it means

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

Solutions

  1. Remove the out argument and use the returned array: idx = jnp.argmax(a)
  2. If a buffer must be filled, do buf = jnp.argmax(a) or buf.at[...].set(...) at the call site
  3. Strip out from kwargs before forwarding to jnp

Example fix

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

Strategy: validation

Validate before calling

kwargs.pop('out', None)  # before forwarding to jnp.argmax

Prevention

When it happens

Trigger: Calling jnp.argmax(a, out=buf), usually by porting NumPy code that reused an output buffer for performance.

Common situations: Copy-pasted NumPy micro-optimized code using out= to avoid allocations; generic wrappers that forward **kwargs including out.

Related errors

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

Appendix: source

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

    smallest index is returned.

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

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

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

@api.jit(static_argnames=('axis', 'keepdims'), inline=True)
def _argmax(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 argmax of an empty sequence")
  # TODO(phawkins): use an int64 index if the dimension is large enough.
  result = lax.argmax(a, _canonicalize_axis(axis, a.ndim), int)
  return expand_dims(result, dims) if keepdims else result

View on GitHub (pinned to 1e1c6a8fc0)