jax-ml/jax · error · NotImplementedError

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

Error message

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

What it means

jnp.nanargmin does not support the out parameter; JAX's functional model returns new arrays rather than mutating caller-supplied buffers.

Source

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

    - :func:`jax.numpy.argmin`: return the index of the minimum value.
    - :func:`jax.numpy.nanargmax`: compute ``argmax`` while ignoring NaN values.

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

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

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


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


@api.jit(static_argnums=(2,))
def _roll_dynamic(a: Array, shift: Array, axis: Sequence[int]) -> Array:
  b_shape = lax.broadcast_shapes(shift.shape, np.shape(axis))
  if len(b_shape) != 1:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Drop out and capture the return value
Pop out from kwargs in shared wrappers before calling jnp
Use functional updates (x = x.at[i].set(...)) for buffer-like semantics

Example fix

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

Defensive patterns

Strategy: validation

Validate before calling

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

Prevention

Remove out= when migrating NumPy nanargmin
Use x.at[...].set(...) for buffer-like updates
Keep jnp calls purely functional

When it happens

Trigger: Calling jnp.nanargmin(a, out=buf), typically from NumPy code migration.

Common situations: Legacy NumPy allocation-avoidance idioms; generic reduction wrappers forwarding **kwargs.

Related errors

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