jax-ml/jax · error · TypeError

Invalid set in op; valid range is

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.

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.

Appendix: 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)