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

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

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.

Appendix: 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)