{"record":{"id":"8c440724576acb53","repo":"jax-ml/jax","slug":"the-out-argument-to-jnp-choose-is-not-supported","errorCode":null,"errorMessage":"The 'out' argument to jnp.choose is not supported.","messagePattern":"The 'out' argument to jnp\\.choose is not supported\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/lax_numpy.py","lineNumber":4963,"sourceCode":"\n    In the more general case, ``choices`` may be a sequence of array-like\n    objects with any broadcast-compatible shapes.\n\n    >>> choice_1 = jnp.array([1, 2, 3, 4])\n    >>> choice_2 = 99\n    >>> choice_3 = jnp.array([[10],\n    ...                       [20],\n    ...                       [30]])\n    >>> a = jnp.array([[0, 1, 2, 0],\n    ...                [1, 2, 0, 1],\n    ...                [2, 0, 1, 2]])\n    >>> jnp.choose(a, [choice_1, choice_2, choice_3], mode='wrap')\n    Array([[ 1, 99, 10,  4],\n           [99, 20,  3, 99],\n           [30,  2, 99, 30]], dtype=int32)\n  \"\"\"\n  if out is not None:\n    raise NotImplementedError(\"The 'out' argument to jnp.choose is not supported.\")\n  a, *choices = util.ensure_arraylike_tuple('choose', (a, *choices))\n  if not issubdtype(a.dtype, np.integer):\n    raise ValueError(\"`a` array must be integer typed\")\n  N = len(choices)\n\n  if mode == 'raise':\n    arr: Array = core.concrete_or_error(asarray, a,\n      \"The error occurred because jnp.choose was jit-compiled\"\n      \" with mode='raise'. Use mode='wrap' or mode='clip' instead.\")\n    if reductions.any((arr < 0) | (arr >= N)):\n      raise ValueError(\"invalid entry in choice array\")\n  elif mode == 'wrap':\n    arr = asarray(a) % N\n  elif mode == 'clip':\n    arr = clip(a, 0, N - 1)\n  else:\n    raise ValueError(f\"mode={mode!r} not understood. Must be 'raise', 'wrap', or 'clip'\")\n","sourceCodeStart":4945,"sourceCodeEnd":4981,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/lax_numpy.py#L4945-L4981","documentation":"jnp.choose mirrors numpy's choose but, like other out= parameters in JAX, cannot write into a preallocated buffer because arrays are immutable; passing out raises NotImplementedError.","triggerScenarios":"jnp.choose(a, choices, out=buf) — code using numpy's out parameter for choose.","commonSituations":"Ported numpy code that preallocates outputs for performance; shared utility functions that thread an out parameter through calls.","solutions":["Remove out= and use the return value","If out semantics are needed, emulate with out = out.at[...].set(jnp.choose(...))"],"exampleFix":"// before\njnp.choose(a, choices, out=out)\n// after\nout = jnp.choose(a, choices)","handlingStrategy":"type-guard","validationCode":"assert out is None","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Remove out= usage when porting numpy; JAX arrays are immutable"],"tags":["jnp-choose","out-argument","immutable-arrays","numpy-parity"],"backgroundTag":"unsupported-out-parameter","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}