jax-ml/jax · error · ValueError

indices must have an integer type

Error message

indices must have an integer type

What it means

The dtype rule of lax.gather requires the indices array to have an integer dtype. Passing float (or other non-integer) indices raises this ValueError before any gathering occurs, since gather offsets must be integral. The operand dtype is returned unchanged when indices are valid.

Source

Thrown at jax/_src/lax/slicing.py:1774

    ur_rule=_dynamic_update_slice_ur_rule)
ad.primitive_jvps[dynamic_update_slice_p] = _dynamic_update_slice_jvp
ad.primitive_transposes[dynamic_update_slice_p] = \
    _dynamic_update_slice_transpose_rule
batching.primitive_batchers[dynamic_update_slice_p] = \
    _dynamic_update_slice_batching_rule

def _dynamic_update_slice_lower(ctx, x, update, *start_indices):
  aval_out, = ctx.avals_out
  out = mlir.dynamic_update_slice(ctx, aval_out, x, update,
                                  start_indices=start_indices)
  return [mlir.lower_with_sharding_in_types(ctx, out, aval_out)]

mlir.register_lowering(dynamic_update_slice_p, _dynamic_update_slice_lower)


def _gather_dtype_rule(operand, indices, *, fill_value, **kwargs):
  if not dtypes.issubdtype(indices.dtype, np.integer):
    raise ValueError("indices must have an integer type")
  return operand.dtype

_rank = lambda arr: len(arr.shape)

def _is_sorted(dims, op_name, name):
  for i in range(1, len(dims)):
    if dims[i] < dims[i - 1]:
      raise TypeError(f"{name} in {op_name} op must be sorted; got {dims}")

def _dims_in_range(dims, rank, op_name, name):
  for dim in dims:
    if dim < 0 or dim >= rank:
      raise TypeError(f"Invalid {name} set in {op_name} op; valid range is "
                      f"[0, {rank}); got: {dim}.")

def _sorted_dims_in_range(dims, rank, op_name, name):
  if len(dims) == 0:
    return

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast indices to an integer dtype: indices = indices.astype(jnp.int32).
2. For sub-pixel / fractional coordinates use interpolation utilities (e.g. jax.scipy.ndimage.map_coordinates) instead of gather.
3. Round before casting if coordinates are fractional: jnp.round(coords).astype(jnp.int32).

Example fix

# before
coords = jnp.array([[1.5, 2.5], [0.1, 3.9]])
out = lax.gather(img, coords, dnums, slice_sizes=(1, 1))  # ValueError

# after
coords = jnp.round(coords).astype(jnp.int32)
out = lax.gather(img, coords, dnums, slice_sizes=(1, 1))

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

def are_integer_indices(indices: jax.Array) -> bool:
    return jnp.issubdtype(indices.dtype, jnp.integer)

When it happens

Trigger: Calling lax.gather(operand, indices, dimension_numbers, slice_sizes, ...) where indices is a float array (e.g. from normalized coordinates, division, or interpolation). Also hit via jnp.take-like code paths that lower to gather with float index tensors.

Common situations: Image resampling / bilinear sampling code computing pixel coordinates in float32; feeding model outputs (float logits) directly as positions; forgetting to cast positions from float32 to int32 after coordinate arithmetic.

Related errors

• {} does not accept dtype {}. Accepted dtypes are subtypes of
• logical reduction requires operand dtype bool or int, got {o
• logaddexp2 requires floating-point or complex inputs; got {x
• index arguments to dynamic_update_slice must be integers of
• {name} in {op_name} op must be sorted; got {dims}

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