jax-ml/jax · error · NotImplementedError

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

Error message

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

What it means

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.

Source

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

    In the more general case, ``choices`` may be a sequence of array-like
    objects with any broadcast-compatible shapes.

    >>> choice_1 = jnp.array([1, 2, 3, 4])
    >>> choice_2 = 99
    >>> choice_3 = jnp.array([[10],
    ...                       [20],
    ...                       [30]])
    >>> a = jnp.array([[0, 1, 2, 0],
    ...                [1, 2, 0, 1],
    ...                [2, 0, 1, 2]])
    >>> jnp.choose(a, [choice_1, choice_2, choice_3], mode='wrap')
    Array([[ 1, 99, 10,  4],
           [99, 20,  3, 99],
           [30,  2, 99, 30]], dtype=int32)
  """
  if out is not None:
    raise NotImplementedError("The 'out' argument to jnp.choose is not supported.")
  a, *choices = util.ensure_arraylike_tuple('choose', (a, *choices))
  if not issubdtype(a.dtype, np.integer):
    raise ValueError("`a` array must be integer typed")
  N = len(choices)

  if mode == 'raise':
    arr: Array = core.concrete_or_error(asarray, a,
      "The error occurred because jnp.choose was jit-compiled"
      " with mode='raise'. Use mode='wrap' or mode='clip' instead.")
    if reductions.any((arr < 0) | (arr >= N)):
      raise ValueError("invalid entry in choice array")
  elif mode == 'wrap':
    arr = asarray(a) % N
  elif mode == 'clip':
    arr = clip(a, 0, N - 1)
  else:
    raise ValueError(f"mode={mode!r} not understood. Must be 'raise', 'wrap', or 'clip'")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Remove out= and use the return value
  2. If out semantics are needed, emulate with out = out.at[...].set(jnp.choose(...))

Example fix

// before
jnp.choose(a, choices, out=out)
// after
out = jnp.choose(a, choices)
Defensive patterns

Strategy: type-guard

Validate before calling

assert out is None

Prevention

When it happens

Trigger: jnp.choose(a, choices, out=buf) — code using numpy's out parameter for choose.

Common situations: Ported numpy code that preallocates outputs for performance; shared utility functions that thread an out parameter through calls.

Related errors

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