jax-ml/jax · error · TypeError

{name} in {op_name} op must not repeat; got: {dims}.

Error message

{name} in {op_name} op must not repeat; got: {dims}.

What it means

Validator used by gather/scatter shape rules that rejects dimension lists containing duplicates: e.g. offset_dims=(1, 1) or update_window_dims=(0, 2, 0). Each dimension may appear at most once in a given list because the maps between input/output axes must be bijective, matching XLA's constraint.

Source

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

    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>`_
  operator and following the outline of the implementation of
  ShapeInference::InferGatherShape in TensorFlow.
  """

  offset_dims = dimension_numbers.offset_dims

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. De-duplicate and re-sort the offending list: dims = tuple(sorted(set(dims))).
  2. Audit the loop/comprehension that builds the dim list for double insertion.
  3. Prefer jnp.take / x.at[idx].set() which never need manual dim lists.

Example fix

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

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

Strategy: validation

Validate before calling

offset_dims = tuple(sorted(set(offset_dims)))
update_window_dims = tuple(sorted(set(update_window_dims)))

Type guard

def has_no_duplicates(dims: tuple) -> bool:
    return len(set(dims)) == len(dims)

Prevention

When it happens

Trigger: Constructing GatherDimensionNumbers or ScatterDimensionNumbers where a dim list repeats a value, then calling lax.gather or lax.scatter/update with those numbers.

Common situations: Programmatic generation of dim lists (nested loops appending indices twice); merging configs from two call sites; off-by-one when slicing a range producing repeated axes (e.g. [i, i] from range closures).

Related errors

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