jax-ml/jax · error · TypeError

{name} in {op_name} op must be sorted; got {dims}

Error message

{name} in {op_name} op must be sorted; got {dims}

What it means

Helper validator used by the gather/scatter shape rules: dimension lists such as offset_dims, collapsed_slice_dims, start_index_map, or update_window_dims must be given in sorted (ascending) order. If any later element is smaller than the previous one, this TypeError names the offending list. JAX requires sorted lists rather than silently sorting them, to match XLA semantics.

Source

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

  aval_out, = ctx.avals_out
  out = mlir.dynamic_update_slice(ctx, aval_out, x, update,
                                  start_indices=start_indices)
  return [mlir.lower_with_sharding_in_types(ctx, out, aval_out)]

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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Sort the listed dims before constructing the dimension numbers: offset_dims=tuple(sorted(offset_dims)).
  2. Prefer higher-level APIs (jnp.take, jnp.take_along_axis, x.at[idx].set(y)) which build dimension numbers for you.
  3. Double-check each list in GatherDimensionNumbers/ScatterDimensionNumbers against the docstring ordering requirement.

Example fix

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

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

Strategy: validation

Validate before calling

offset_dims = tuple(sorted(offset_dims))
collapsed_slice_dims = tuple(sorted(collapsed_slice_dims))
start_index_map = tuple(sorted(start_index_map))
dnums = lax.GatherDimensionNumbers(offset_dims, collapsed_slice_dims, start_index_map)

Type guard

def dims_are_sorted(dims: tuple) -> bool:
    return all(dims[i] >= dims[i-1] for i in range(1, len(dims)))

Prevention

When it happens

Trigger: Building lax.GatherDimensionNumbers or lax.ScatterDimensionNumbers by hand with an unsorted tuple, e.g. offset_dims=(2, 0) or start_index_map=(3, 1), then calling lax.gather/lax.scatter with those dimension numbers.

Common situations: Hand-rolling embedding lookups or scatter updates instead of jnp.take / .at[]; porting XLA or TensorFlow gather semantics where dims were listed in a different order; dynamically generating dimension numbers with a loop that appends out of order.

Related errors

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