jax-ml/jax · error · NotImplementedError

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

Error message

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

What it means

JAX arrays are immutable, so nanargmax cannot write into a caller-supplied out array. The wrapper rejects any non-None out argument to jnp.nanargmax.

Solutions

Remove out and use the return value
Assign afterwards with .at[...].set(...) if a preallocated structure must be populated
Do not forward out to jnp functions

Example fix

// before
jnp.nanargmax(a, out=buf)
// after
buf = jnp.nanargmax(a)

Defensive patterns

Strategy: validation

Validate before calling

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

Prevention

Drop out= from jnp calls
Assign results functionally
Sanitize kwargs in shared reduction helpers

When it happens

Trigger: Calling jnp.nanargmax(a, out=buf) — almost always ported NumPy code or a kwargs-forwarding wrapper.

Common situations: Copy-paste from NumPy pipelines that used out=; refactored shared reduction helpers that pass out uniformly.

Related errors

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

Appendix: source

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

    >>> jnp.argmax(x)
    Array(4, dtype=int32)

    Using ``nanargmax`` returns the index of the maximum non-NaN value.

    >>> jnp.nanargmax(x)
    Array(2, dtype=int32)

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

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


@api.jit(static_argnames=('axis', 'keepdims'))
def _nanargmax(a: Array, axis: int | None = None, keepdims: bool = False):
  if not issubdtype(a.dtype, np.inexact):
    return argmax(a, axis=axis, keepdims=keepdims)
  nan_mask = ufuncs.isnan(a)
  a = where(nan_mask, -np.inf, a)
  res = argmax(a, axis=axis, keepdims=keepdims)
  return where(reductions.all(nan_mask, axis=axis, keepdims=keepdims), -1, res)


@export
def nanargmin(
    a: ArrayLike,
    axis: int | None = None,

View on GitHub (pinned to 1e1c6a8fc0)