{"record":{"id":"31b6a955dc29950f","repo":"jax-ml/jax","slug":"the-out-argument-to-jnp-nanargmax-is-not-support","errorCode":null,"errorMessage":"The 'out' argument to jnp.nanargmax is not supported.","messagePattern":"The 'out' argument to jnp\\.nanargmax is not supported\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/lax_numpy.py","lineNumber":8406,"sourceCode":"    >>> jnp.argmax(x)\n    Array(4, dtype=int32)\n\n    Using ``nanargmax`` returns the index of the maximum non-NaN value.\n\n    >>> jnp.nanargmax(x)\n    Array(2, dtype=int32)\n\n    >>> x = jnp.array([[1, 3, jnp.nan],\n    ...                [5, 4, jnp.nan]])\n    >>> jnp.nanargmax(x, axis=1)\n    Array([1, 0], dtype=int32)\n\n    >>> jnp.nanargmax(x, axis=1, keepdims=True)\n    Array([[1],\n           [0]], dtype=int32)\n  \"\"\"\n  if out is not None:\n    raise NotImplementedError(\"The 'out' argument to jnp.nanargmax is not supported.\")\n  a = util.ensure_arraylike(\"nanargmax\", a)\n  return _nanargmax(a, None if axis is None else operator.index(axis), keepdims=bool(keepdims))\n\n\n@api.jit(static_argnames=('axis', 'keepdims'))\ndef _nanargmax(a: Array, axis: int | None = None, keepdims: bool = False):\n  if not issubdtype(a.dtype, np.inexact):\n    return argmax(a, axis=axis, keepdims=keepdims)\n  nan_mask = ufuncs.isnan(a)\n  a = where(nan_mask, -np.inf, a)\n  res = argmax(a, axis=axis, keepdims=keepdims)\n  return where(reductions.all(nan_mask, axis=axis, keepdims=keepdims), -1, res)\n\n\n@export\ndef nanargmin(\n    a: ArrayLike,\n    axis: int | None = None,","sourceCodeStart":8388,"sourceCodeEnd":8424,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/lax_numpy.py#L8388-L8424","documentation":"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.","triggerScenarios":"Calling jnp.nanargmax(a, out=buf) — almost always ported NumPy code or a kwargs-forwarding wrapper.","commonSituations":"Copy-paste from NumPy pipelines that used out=; refactored shared reduction helpers that pass out uniformly.","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"],"exampleFix":"// before\njnp.nanargmax(a, out=buf)\n// after\nbuf = jnp.nanargmax(a)\n","handlingStrategy":"validation","validationCode":"kwargs.pop('out', None)  # before calling jnp.nanargmax","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Drop out= from jnp calls","Assign results functionally","Sanitize kwargs in shared reduction helpers"],"tags":["jax","nanargmax","unsupported-argument","numpy-compat"],"backgroundTag":"unsupported-out-parameter","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}