{"record":{"id":"bfa24c105a6c4f4e","repo":"jax-ml/jax","slug":"the-out-argument-to-jnp-argmax-is-not-supported","errorCode":null,"errorMessage":"The 'out' argument to jnp.argmax is not supported.","messagePattern":"The 'out' argument to jnp\\.argmax is not supported\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/lax_numpy.py","lineNumber":8273,"sourceCode":"    smallest index is returned.\n\n  Examples:\n    >>> x = jnp.array([1, 3, 5, 4, 2])\n    >>> jnp.argmax(x)\n    Array(2, dtype=int32)\n\n    >>> x = jnp.array([[1, 3, 2],\n    ...                [5, 4, 1]])\n    >>> jnp.argmax(x, axis=1)\n    Array([1, 0], dtype=int32)\n\n    >>> jnp.argmax(x, axis=1, keepdims=True)\n    Array([[1],\n           [0]], dtype=int32)\n  \"\"\"\n  arr = util.ensure_arraylike(\"argmax\", a)\n  if out is not None:\n    raise NotImplementedError(\"The 'out' argument to jnp.argmax is not supported.\")\n  return _argmax(arr, None if axis is None else operator.index(axis),\n                 keepdims=bool(keepdims))\n\n@api.jit(static_argnames=('axis', 'keepdims'), inline=True)\ndef _argmax(a: Array, axis: int | None = None, keepdims: bool = False) -> Array:\n  if axis is None:\n    dims = list(range(np.ndim(a)))\n    a = ravel(a)\n    axis = 0\n  else:\n    dims = [axis]\n  if a.shape[axis] == 0:\n    raise ValueError(\"attempt to get argmax of an empty sequence\")\n  # TODO(phawkins): use an int64 index if the dimension is large enough.\n  result = lax.argmax(a, _canonicalize_axis(axis, a.ndim), int)\n  return expand_dims(result, dims) if keepdims else result\n\n","sourceCodeStart":8255,"sourceCodeEnd":8291,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/lax_numpy.py#L8255-L8291","documentation":"jnp.argmax does not support the out parameter because JAX arrays are immutable and jit-traced functions cannot write into caller-provided buffers, unlike NumPy.","triggerScenarios":"Calling jnp.argmax(a, out=buf), usually by porting NumPy code that reused an output buffer for performance.","commonSituations":"Copy-pasted NumPy micro-optimized code using out= to avoid allocations; generic wrappers that forward **kwargs including out.","solutions":["Remove the out argument and use the returned array: idx = jnp.argmax(a)","If a buffer must be filled, do buf = jnp.argmax(a) or buf.at[...].set(...) at the call site","Strip out from kwargs before forwarding to jnp"],"exampleFix":"// before\njnp.argmax(a, axis=0, out=buf)\n// after\nbuf = jnp.argmax(a, axis=0)\n","handlingStrategy":"validation","validationCode":"kwargs.pop('out', None)  # before forwarding to jnp.argmax","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Never pass out= to jnp reductions","Strip out in generic wrappers","Use return-value assignment instead of buffers"],"tags":["jax","argmax","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"}