jax-ml/jax · error · TypeError

Invalid {name} set in {op_name} op; valid range is [0, {rank

Error message

Invalid {name} set in {op_name} op; valid range is [0, {rank}); got: {invalid_dim}.

What it means

Optimized bounds check used by gather/scatter shape rules for already-sorted dim lists: it only inspects the first and last elements, and raises if dims[0] < 0 or dims[-1] >= rank. It produces the same 'Invalid {name} set' message as the linear check but identifies the boundary offender; empty lists pass trivially.

Source

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

    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
  invalid_dim = None
  if dims[0] < 0:
    invalid_dim = dims[0]
  elif dims[-1] >= rank:
    invalid_dim = dims[-1]
  if invalid_dim:
    raise TypeError(f"Invalid {name} set in {op_name} op; valid range is "
                    f"[0, {rank}); got: {invalid_dim}.")

def _no_duplicate_dims(dims, op_name, name):
  if len(set(dims)) != len(dims):
    raise TypeError(f"{name} in {op_name} op must not repeat; got: {dims}.")

def _disjoint_dims(dims1, dims2, op_name, name1, name2):
  if not set(dims1).isdisjoint(set(dims2)):
    raise TypeError(f"{name1} and {name2} in {op_name} op must be disjoint; "
                    f"got: {dims1} and {dims2}.")

def _gather_shape_rule(operand, indices, *, dimension_numbers,
                       slice_sizes, unique_indices, indices_are_sorted,
                       mode, fill_value):
  """Validates the well-formedness of the arguments to Gather.

  The code implements the checks based on the detailed operation semantics of
  XLA's `Gather <https://www.openxla.org/xla/operation_semantics#gather>`_

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Check dims[0] >= 0 and dims[-1] < rank for every list in the dimension numbers; fix the value shown in the message.
Derive dim lists from operand.ndim at runtime rather than constants.
Use jnp.take / .at[] style APIs to avoid manual dimension numbers entirely.

Example fix

# before
rank = x.ndim  # 3
dnums = lax.GatherDimensionNumbers(
    offset_dims=(), collapsed_slice_dims=(0, 3), start_index_map=(0,))
out = lax.gather(x, idx, dnums, slice_sizes=(1,))  # 3 >= rank -> TypeError

# after
dnums = lax.GatherDimensionNumbers(
    offset_dims=(1,), collapsed_slice_dims=(0, 2), start_index_map=(0,))
out = lax.gather(x, idx, dnums, slice_sizes=(1, 1))

Defensive patterns

Strategy: validation

Validate before calling

def validate_sorted_dims(dims, rank, name):
    if dims and (dims[0] < 0 or dims[-1] >= rank):
        raise ValueError(f"{name} out of range [0, {rank}): {dims}")
    return dims

Prevention

Validate list extremes against operand rank when dim lists are already sorted.
Regenerate dimension numbers automatically when operand rank changes in a refactor.
Keep a single helper that builds dnums so rank checks happen in one place.

When it happens

Trigger: Supplying a sorted dim list whose extremes are out of bounds, e.g. start_index_map=(0, 5) for a rank-3 operand to lax.gather, or update_window_dims starting at a negative value to lax.scatter.

Common situations: Same class of bugs as the unsorted variant: dimension numbers copied from another rank's config, computing 'last axis' as rank instead of rank-1, or dims derived from index tensors whose depth exceeds the operand rank.

Related errors

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