{"record":{"id":"6e5967a92c4dfe7e","repo":"jax-ml/jax","slug":"invalid-entry-in-choice-array","errorCode":null,"errorMessage":"invalid entry in choice array","messagePattern":"invalid entry in choice array","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/lax_numpy.py","lineNumber":4974,"sourceCode":"    ...                [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\n  arr, *choices = broadcast_arrays(arr, *choices)\n  return array(choices)[(arr,) + indices(arr.shape, sparse=True)]\n\n\ndef _atleast_nd(x: ArrayLike, n: int) -> Array:\n  m = np.ndim(x)\n  return lax.broadcast(x, (1,) * (n - m)) if m < n else asarray(x)\n\ndef _block(xs: ArrayLike | list[Any]) -> tuple[Array, int]:\n  if isinstance(xs, tuple):\n    raise ValueError(\"jax.numpy.block does not allow tuples, got {}\"","sourceCodeStart":4956,"sourceCodeEnd":4992,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/lax_numpy.py#L4956-L4992","documentation":"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.","triggerScenarios":"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.","commonSituations":"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'.","solutions":["Clip or wrap indices explicitly: jnp.clip(a, 0, len(choices)-1) or a % len(choices)","Pass mode='wrap' or mode='clip' to jnp.choose to get NumPy-like out-of-range handling","Fix the upstream index computation (off-by-one, wrong argmax axis)"],"exampleFix":"// before\nout = jnp.choose(idx, [c0, c1, c2])  # idx may contain 3+\n// after\nout = jnp.choose(jnp.clip(idx, 0, 2), [c0, c1, c2])\n// or\nout = jnp.choose(idx, [c0, c1, c2], mode='clip')","handlingStrategy":"validation","validationCode":"N = len(choices)\nif jnp.any((a < 0) | (a >= N)):\n    a = jnp.clip(a, 0, N - 1)  # or a % N","typeGuard":null,"tryCatchPattern":"try:\n    out = jnp.choose(a, choices)\nexcept ValueError:\n    out = jnp.choose(jnp.clip(a, 0, len(choices) - 1), choices)","preventionTips":["Use mode='clip' or mode='wrap' when indices are not guaranteed in range","Never use mode='raise' inside jax.jit"],"tags":["jax","choose","index-out-of-range","valueerror"],"backgroundTag":"index-out-of-bounds","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}