jax-ml/jax · error · ValueError

mode={mode!r} not understood. Must be 'raise', 'wrap', or 'c

Error message

mode={mode!r} not understood. Must be 'raise', 'wrap', or 'clip'

What it means

jnp.choose only accepts mode in {'raise', 'wrap', 'clip'}. Any other string (or non-string) value for the mode keyword raises this ValueError with the offending value echoed back. This matches NumPy's choose modes but JAX adds tracer-related constraints on 'raise' under jit.

Source

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

  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 {}"
                     .format(xs))
  elif isinstance(xs, list):
    if len(xs) == 0:
      raise ValueError("jax.numpy.block does not allow empty list arguments")
    xs_tup, depths = unzip2([_block(x) for x in xs])
    if any(d != depths[0] for d in depths[1:]):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use exactly one of 'raise', 'wrap', or 'clip' (lowercase)
2. If under jit, prefer 'wrap' or 'clip' since 'raise' needs concrete values

Example fix

// before
out = jnp.choose(a, ch, mode='fill')
// after
out = jnp.choose(a, ch, mode='clip')

assert mode in ('raise', 'wrap', 'clip'), f"bad mode: {mode}"

def is_valid_choose_mode(mode) -> bool:
    return mode in ('raise', 'wrap', 'clip')

When it happens

Trigger: Calling jnp.choose(a, choices, mode='Raise') (wrong case), mode='fill' (not a real mode), or passing mode=None. Also typos like mode='wrapp'.

Common situations: Copy-paste from code that used np.take (which uses mode='wrap'/'clip' too but also 'raise' semantics differently), or assumption that NumPy's error message strings differ.

Related errors

• `a` array must be integer typed
• invalid entry in choice array
• 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/5c9eb9d26150356a. Report an issue: GitHub.