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

jnp.nanargmax does not support the out parameter for the same reason as other JAX functions: arrays are immutable and results are returned, not written into buffers.

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)

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.