{"record":{"id":"b2e277a31925c399","repo":"jax-ml/jax","slug":"gather-index-leaf-dimension-must-be-within-0-ran","errorCode":null,"errorMessage":"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}.","messagePattern":"Gather index leaf dimension must be within \\[0, rank\\(indices\\) \\+ 1\\)\\. rank\\(indices\\) is (.+?) and gather index leaf dimension is (.+?)\\.","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/slicing.py","lineNumber":1835,"sourceCode":"  XLA's `Gather <https://www.openxla.org/xla/operation_semantics#gather>`_\n  operator and following the outline of the implementation of\n  ShapeInference::InferGatherShape in TensorFlow.\n  \"\"\"\n\n  offset_dims = dimension_numbers.offset_dims\n  collapsed_slice_dims = dimension_numbers.collapsed_slice_dims\n  operand_batching_dims = dimension_numbers.operand_batching_dims\n  start_indices_batching_dims = dimension_numbers.start_indices_batching_dims\n  start_index_map = dimension_numbers.start_index_map\n\n  # Note: in JAX, index_vector_dim is always computed as below, cf. the\n  # documentation of the GatherDimensionNumbers class.\n  index_vector_dim = _rank(indices) - 1\n\n  # This case should never happen in JAX, due to the implicit construction of\n  # index_vector_dim, but is included for completeness.\n  if _rank(indices) < index_vector_dim or index_vector_dim < 0:\n    raise TypeError(f\"Gather index leaf dimension must be within [0, rank(\"\n                    f\"indices) + 1). rank(indices) is {_rank(indices)} and \"\n                    f\"gather index leaf dimension is {index_vector_dim}.\")\n\n  # Start ValidateGatherDimensions\n  # In the error messages output by XLA, \"offset_dims\" is called \"Output window\n  # dimensions\" in error messages. For consistency's sake, our error messages\n  # stick to \"offset_dims\".\n  _is_sorted(offset_dims, \"gather\", \"offset_dims\")\n  _no_duplicate_dims(offset_dims, \"gather\", \"offset_dims\")\n\n  output_offset_dim_count = len(offset_dims)\n  output_shape_rank = len(offset_dims) + _rank(indices) - 1\n\n  for i in range(output_offset_dim_count):\n    offset_dim = offset_dims[i]\n    if offset_dim < 0 or offset_dim >= output_shape_rank:\n      raise TypeError(f\"Offset dimension {i} in gather op is out of bounds; \"\n                      f\"got {offset_dim}, but should have been in \"","sourceCodeStart":1817,"sourceCodeEnd":1853,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/slicing.py#L1817-L1853","documentation":"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)].","triggerScenarios":"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.","commonSituations":"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.","solutions":["Give indices at least one dimension, e.g. indices = idx[None, :] or idx.reshape(1, -1) so the trailing axis holds index vectors.","Guard before calling: if indices.ndim == 0: indices = indices[None].","For a single index use jnp.take(operand, int(idx)) instead of lax.gather."],"exampleFix":"# before\nidx = jnp.array(3)  # scalar\nout = lax.gather(x, idx, dnums, slice_sizes=(1,))  # index_vector_dim = -1\n\n# after\nidx = jnp.array([3])  # shape (1,)\nout = lax.gather(x, idx, dnums, slice_sizes=(1,))","handlingStrategy":"type-guard","validationCode":"if indices.ndim == 0:\n    indices = indices[None]  # shape () -> (1,)","typeGuard":"def valid_gather_indices(indices: jax.Array) -> bool:\n    return indices.ndim >= 1","tryCatchPattern":null,"preventionTips":["Never pass scalar indices to lax.gather; reshape to at least 1-D (idx[None]).","Use jnp.take for single-index lookups.","Guard batched loops where squeezing can collapse index tensors to rank 0."],"tags":["jax","lax","gather","index-vector-dim"],"backgroundTag":"invalid-index-shape","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}