What causes "`a` array must be integer typed" in jax-ml/jax?

Cast the index array to an integer type before calling choose: jnp.choose(a.astype(jnp.int32), choices)

jax-ml/jax · error · ValueError

`a` array must be integer typed

Error message

`a` array must be integer typed

What it means

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.

Source

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

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

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

View on GitHub (pinned to 1e1c6a8fc0)

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))

Example fix

// before
result = jnp.choose(scores, choices)  # scores is float32
// after
result = jnp.choose(scores.astype(jnp.int32), choices)

Defensive patterns

Strategy: validation

Validate before calling

a = jnp.asarray(a)
if not jnp.issubdtype(a.dtype, jnp.integer):
    a = a.astype(jnp.int32)

Type guard

def is_integer_index_array(a) -> bool:
    return jnp.issubdtype(jnp.asarray(a).dtype, jnp.integer)

Try / catch

try:
    out = jnp.choose(a, choices)
except ValueError as e:
    if 'must be integer typed' in str(e):
        out = jnp.choose(jnp.asarray(a).astype(jnp.int32), choices)
    else:
        raise

Prevention

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

When it happens

Trigger: 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.

Common situations: 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.

Related errors

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