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: {dim}.

What it means

Linear-scan validator (used by gather/scatter shape rules) that rejects dimension numbers outside [0, rank): any dim < 0 or >= rank in lists like offset_dims, collapsed_slice_dims, start_index_map, gathered_from_slice_dims, or scatter_dims raises this TypeError showing the offending dim. It fires on unsorted lists too, since it checks every element individually.

Source

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

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
  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}.")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Print operand rank (operand.ndim) and clamp all listed dims to 0..rank-1; fix the out-of-range dim shown in the message.
Compute dimension numbers relative to rank at runtime instead of hardcoding, e.g. collapsed_slice_dims=(operand.ndim - 1,).
Replace hand-written dnums with jnp.take / x.at[] helpers when possible.

Example fix

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

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

Defensive patterns

Strategy: validation

Validate before calling

def check_dims(dims, rank, name):
    bad = [d for d in dims if d < 0 or d >= rank]
    assert not bad, f"{name} out of range [0, {rank}): {bad}"
    return tuple(sorted(set(dims)))

Prevention

Derive dim lists from operand.ndim at runtime instead of hardcoding constants.
Remember the last valid axis is rank-1, not rank.
Run check_dims on every list before constructing GatherDimensionNumbers/ScatterDimensionNumbers.

When it happens

Trigger: Passing dimension numbers that reference an axis beyond the array rank, e.g. collapsed_slice_dims=(3,) for a rank-2 operand to lax.gather, or negative dims like update_window_dims=(-1,) to lax.scatter.

Common situations: Reusing dimension numbers written for a different rank (e.g. ported from a batched version); off-by-one when computing the last axis as rank instead of rank-1; changing operand rank during a refactor without updating hand-built GatherDimensionNumbers/ScatterDimensionNumbers.

Related errors

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