jax-ml/jax · error · ValueError

invalid entry in choice array

Error message

invalid entry in choice array

What it means

With mode='raise' (the default), jnp.choose validates that every element of the index array `a` lies in [0, N) where N is the number of choice arrays. Any out-of-range value (negative or >= N) triggers this ValueError. This is the JAX equivalent of NumPy's 'invalid entry in choice array' error.

Source

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

    ...                [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'")

  arr, *choices = broadcast_arrays(arr, *choices)
  return array(choices)[(arr,) + indices(arr.shape, sparse=True)]


def _atleast_nd(x: ArrayLike, n: int) -> Array:
  m = np.ndim(x)
  return lax.broadcast(x, (1,) * (n - m)) if m < n else asarray(x)

def _block(xs: ArrayLike | list[Any]) -> tuple[Array, int]:
  if isinstance(xs, tuple):
    raise ValueError("jax.numpy.block does not allow tuples, got {}"

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Clip or wrap indices explicitly: jnp.clip(a, 0, len(choices)-1) or a % len(choices)
2. Pass mode='wrap' or mode='clip' to jnp.choose to get NumPy-like out-of-range handling
3. Fix the upstream index computation (off-by-one, wrong argmax axis)

Example fix

// before
out = jnp.choose(idx, [c0, c1, c2])  # idx may contain 3+
// after
out = jnp.choose(jnp.clip(idx, 0, 2), [c0, c1, c2])
// or
out = jnp.choose(idx, [c0, c1, c2], mode='clip')

N = len(choices)
if jnp.any((a < 0) | (a >= N)):
    a = jnp.clip(a, 0, N - 1)  # or a % N

try:
    out = jnp.choose(a, choices)
except ValueError:
    out = jnp.choose(jnp.clip(a, 0, len(choices) - 1), choices)

When it happens

Trigger: jnp.choose(a, choices) with mode='raise' where a contains e.g. 5 but only 3 choices are given, or negative values from computations like a - 1. Also raised under jit only as a concrete-value error via core.concrete_or_error when values are not concrete.

Common situations: Index arrays derived from argmax over a different-sized axis than the choices list; off-by-one errors where indices are 1-based but choices are 0-based; using mode='raise' inside jit which first fails with a tracer error telling you to use 'wrap' or 'clip'.

Related errors

• `a` array must be integer typed
• mode={mode!r} not understood. Must be 'raise', 'wrap', or 'c
• Invalid value "{default}" for JAX flag {name}
• new enum value must be in {enum_values}, got {new_val} of ty
• new enum value must be None or in {enum_values}, got {new_va

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