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.

Solutions

  1. Cast the index array to an integer type before calling choose: jnp.choose(a.astype(jnp.int32), choices)
  2. Regenerate the index array with an integer-producing API (e.g. jax.random.randint or argmax)
  3. 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

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.

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