jax-ml/jax · error · TypeError

and in op must be disjoint; got: and .

Error message

{name1} and {name2} in {op_name} op must be disjoint; got: {dims1} and {dims2}.

What it means

Validator used by gather/scatter shape rules that requires two dimension lists to be pairwise disjoint, e.g. collapsed_slice_dims and offset_dims in gather, or update_window_dims and inserted_window_dims in scatter. Sharing any axis between the two lists raises this TypeError showing both lists, because an axis cannot play two roles in the dimension mapping.

Solutions

1. Partition the rank: every axis 0..rank-1 goes into exactly one of the two lists (e.g. offset_dims + collapsed_slice_dims should cover the slice dims without overlap).
2. Remove the shared axis from one list, usually from offset_dims / update_window_dims.
3. Validate with set(list1).isdisjoint(list2) before calling, or switch to jnp.take / .at[] helpers.

Example fix

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

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

assert set(collapsed_slice_dims).isdisjoint(offset_dims), \
    f"overlapping dims: {collapsed_slice_dims} vs {offset_dims}"
dnums = lax.GatherDimensionNumbers(offset_dims, collapsed_slice_dims, start_index_map)

When it happens

Trigger: Building lax.GatherDimensionNumbers with an axis present in both collapsed_slice_dims and offset_dims; or lax.ScatterDimensionNumbers where an axis appears in both update_window_dims and inserted_window_dims; then calling lax.gather / lax.scatter.

Common situations: Hand-writing dimension numbers to mimic an embedding lookup and double-assigning the axis that holds the index; configs copied between gather and scatter with different disjointness rules; computing one list as 'all dims' and the other non-empty.

Related errors

• in op must be sorted; got
• Invalid set in op; valid range is
• Invalid set in op; valid range is
• in op must not repeat; got: .
• Gather index leaf dimension must be within

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

Appendix: source

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

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
  collapsed_slice_dims = dimension_numbers.collapsed_slice_dims
  operand_batching_dims = dimension_numbers.operand_batching_dims
  start_indices_batching_dims = dimension_numbers.start_indices_batching_dims
  start_index_map = dimension_numbers.start_index_map

View on GitHub (pinned to 1e1c6a8fc0)