jax-ml/jax · error · TypeError

Gather index leaf dimension must be within [0, rank(indices)

Error message

Gather index leaf dimension must be within [0, rank(indices) + 1). rank(indices) is {_rank(indices)} and gather index leaf dimension is {index_vector_dim}.

What it means

Completeness check in the gather shape rule for the index_vector_dim: since JAX implicitly sets index_vector_dim = rank(indices) - 1, it fires only if rank(indices) < index_vector_dim or index_vector_dim < 0, i.e. essentially when indices is rank-0 and the computed leaf dim becomes -1. It mirrors XLA's requirement that the index vector dimension lie in [0, rank(indices)].

Source

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

  XLA's `Gather <https://www.openxla.org/xla/operation_semantics#gather>`_
  operator and following the outline of the implementation of
  ShapeInference::InferGatherShape in TensorFlow.
  """

  offset_dims = dimension_numbers.offset_dims
  collapsed_slice_dims = dimension_numbers.collapsed_slice_dims
  operand_batching_dims = dimension_numbers.operand_batching_dims
  start_indices_batching_dims = dimension_numbers.start_indices_batching_dims
  start_index_map = dimension_numbers.start_index_map

  # Note: in JAX, index_vector_dim is always computed as below, cf. the
  # documentation of the GatherDimensionNumbers class.
  index_vector_dim = _rank(indices) - 1

  # This case should never happen in JAX, due to the implicit construction of
  # index_vector_dim, but is included for completeness.
  if _rank(indices) < index_vector_dim or index_vector_dim < 0:
    raise TypeError(f"Gather index leaf dimension must be within [0, rank("
                    f"indices) + 1). rank(indices) is {_rank(indices)} and "
                    f"gather index leaf dimension is {index_vector_dim}.")

  # Start ValidateGatherDimensions
  # In the error messages output by XLA, "offset_dims" is called "Output window
  # dimensions" in error messages. For consistency's sake, our error messages
  # stick to "offset_dims".
  _is_sorted(offset_dims, "gather", "offset_dims")
  _no_duplicate_dims(offset_dims, "gather", "offset_dims")

  output_offset_dim_count = len(offset_dims)
  output_shape_rank = len(offset_dims) + _rank(indices) - 1

  for i in range(output_offset_dim_count):
    offset_dim = offset_dims[i]
    if offset_dim < 0 or offset_dim >= output_shape_rank:
      raise TypeError(f"Offset dimension {i} in gather op is out of bounds; "
                      f"got {offset_dim}, but should have been in "

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Give indices at least one dimension, e.g. indices = idx[None, :] or idx.reshape(1, -1) so the trailing axis holds index vectors.
2. Guard before calling: if indices.ndim == 0: indices = indices[None].
3. For a single index use jnp.take(operand, int(idx)) instead of lax.gather.

Example fix

# before
idx = jnp.array(3)  # scalar
out = lax.gather(x, idx, dnums, slice_sizes=(1,))  # index_vector_dim = -1

# after
idx = jnp.array([3])  # shape (1,)
out = lax.gather(x, idx, dnums, slice_sizes=(1,))

if indices.ndim == 0:
    indices = indices[None]  # shape () -> (1,)

def valid_gather_indices(indices: jax.Array) -> bool:
    return indices.ndim >= 1

When it happens

Trigger: Passing a 0-D (scalar) indices array to lax.gather, making index_vector_dim = -1; only reachable in JAX when the implicit computation yields a negative leaf dimension.

Common situations: Degenerate index inputs (single scalar index) instead of shape (1,) or (n, k) index arrays; edge cases in batched loops where an empty or squeezed index tensor collapses to a scalar; direct use of the gather primitive with malformed shapes.

Related errors

• indices must have an integer type
• {name} in {op_name} op must be sorted; got {dims}
• Invalid {name} set in {op_name} op; valid range is [0, {rank
• Invalid {name} set in {op_name} op; valid range is [0, {rank
• {name} in {op_name} op must not repeat; got: {dims}.

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