{"record":{"id":"fb0e694cd1d714df","repo":"jax-ml/jax","slug":"a-array-must-be-integer-typed","errorCode":null,"errorMessage":"`a` array must be integer typed","messagePattern":"`a` array must be integer typed","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/lax_numpy.py","lineNumber":4966,"sourceCode":"\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\n  arr, *choices = broadcast_arrays(arr, *choices)\n  return array(choices)[(arr,) + indices(arr.shape, sparse=True)]\n","sourceCodeStart":4948,"sourceCodeEnd":4984,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/lax_numpy.py#L4948-L4984","documentation":"jax.numpy.choose requires the index array `a` to have an integer dtype, because its values are used to select among the choice arrays. If `a` is float or another non-integer type after being converted to an array, jnp.choose raises this ValueError immediately. This mirrors NumPy's choose, which also demands integer indices.","triggerScenarios":"Calling jnp.choose(a, choices) where `a` is a float array (e.g. output of a softmax/argmax-like computation cast to float, or a Python list of floats), or passing a boolean/complex index array. `util.ensure_arraylike_tuple` converts `a`, then `issubdtype(a.dtype, np.integer)` fails.","commonSituations":"Indices produced by jnp.argmax are fine (int), but users often normalize or cast indices to float32 for downstream math and then reuse them in jnp.choose; or they build `a` from np.random.rand (floats) instead of randint.","solutions":["Cast the index array to an integer type before calling choose: jnp.choose(a.astype(jnp.int32), choices)","Regenerate the index array with an integer-producing API (e.g. jax.random.randint or argmax)","If indices are computed in float, verify they are integral values before casting: assert jnp.all(a == jnp.round(a))"],"exampleFix":"// before\nresult = jnp.choose(scores, choices)  # scores is float32\n// after\nresult = jnp.choose(scores.astype(jnp.int32), choices)","handlingStrategy":"validation","validationCode":"a = jnp.asarray(a)\nif not jnp.issubdtype(a.dtype, jnp.integer):\n    a = a.astype(jnp.int32)","typeGuard":"def is_integer_index_array(a) -> bool:\n    return jnp.issubdtype(jnp.asarray(a).dtype, jnp.integer)","tryCatchPattern":"try:\n    out = jnp.choose(a, choices)\nexcept ValueError as e:\n    if 'must be integer typed' in str(e):\n        out = jnp.choose(jnp.asarray(a).astype(jnp.int32), choices)\n    else:\n        raise","preventionTips":["Always cast index arrays to int32/int64 before choose","Keep index tensors in integer dtype through the pipeline; cast to float only at the final math step"],"tags":["jax","numpy","dtype","valueerror","choose"],"backgroundTag":"array-dtype-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}