jax-ml/jax · error · ValueError

{op_name} sources and destinations must be unique, got {}.

Error message

{op_name} sources and destinations must be unique, got {}.

What it means

ppermute (and psend/precv) lowering builds source/destination pair lists; each replica may appear only once as a source and once as a destination, otherwise the permutation is ambiguous. Duplicates are rejected after normalizing indices modulo group size.

Source

Thrown at jax/_src/lax/parallel.py:1174

  partial(_batched_reduction_collective, pmax_p, lambda v, axis_size: v)


pmin_p = core.Primitive('pmin')
pmin_p.def_impl(partial(_allreduce_impl, pmin_p, lax.reduce_min))
pmin_p.def_effectful_abstract_eval(partial(_pmin_pmax_abstract_eval, 'pmin'))
mlir.register_lowering(
    pmin_p, partial(_all_reduce_lowering, lax.min_p, lax.reduce_min))
batching.fancy_primitive_batchers[pmin_p] = \
  partial(_batched_reduction_collective, pmin_p, lambda v, axis_size: v)


def _pcollectives_lowering_common(ctx, *, axis_name, perm, op_name):
  replica_groups = _replica_groups(ctx.module_context.axis_context, axis_name, None)
  group_size = len(replica_groups[0])
  srcs, dsts = unzip2((src % group_size, dst % group_size) for src, dst in perm)
  if not (len(srcs) == len(set(srcs)) and len(dsts) == len(set(dsts))):
    msg = f"{op_name} sources and destinations must be unique, got {{}}."
    raise ValueError(msg.format(perm))

  full_perm = np.zeros((len(replica_groups), len(perm), 2), np.int64)
  for i, grp in enumerate(replica_groups):
    sorted_grp = tuple(sorted(grp))
    if config.raise_on_ppermute_sort_diff.value and sorted_grp != grp:
      raise RuntimeError(
          "Make sure that the axis_name passed to jax.lax.ppermute is in the"
          " same order as the axis_names declared on the mesh. If you want to"
          " allow different order, you can disable the check via `with"
          " jax.raise_on_ppermute_sort_diff(False):` context manager.")
    for j, (src, dst) in enumerate(perm):
      full_perm[i, j, 0] = grp[src]
      full_perm[i, j, 1] = grp[dst]
  full_perm = full_perm.reshape((-1, 2))

  axis_context = ctx.module_context.axis_context
  if isinstance(axis_context, SPMDAxisContext) and axis_context.manual_axes:
    other_args: dict[str, Any] = dict(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Rewrite perm so sources are a permutation of 0..n-1 and destinations likewise
2. Build perm programmatically (e.g. [(i,(i+1)%n) for i in range(n)]) and assert uniqueness in a unit test
3. Check for accidental wrap-around pairs like (0, n) that normalize to (0, 0)

Example fix

// before
perm = [(0,1),(0,2)]
// after
n = 4; perm = [(i,(i+1)%n) for i in range(n)]

def check_perm(perm, n):
    srcs = [s % n for s, _ in perm]; dsts = [d % n for _, d in perm]
    assert len(srcs) == len(set(srcs)) and len(dsts) == len(set(dsts)), 'perm srcs/dsts must be unique'

When it happens

Trigger: Passing a perm to lax.ppermute where the same source (or destination) index appears twice, e.g. [(0,1),(0,2)].

Common situations: Hand-writing permutation tables; accidentally including an identity pair (i,i) plus another pair using i; off-by-one in modulo group size math.

Related errors

• `perm` passed to `jax.lax.ppermute` must be a list or a tupl
• Replica groups must be equally sized
• {name} only accepts inputs that are unreduced. Got {aval.str
• {full_name} must be a pytree prefix with bool leaves or a tu
• unbound axis name: {axis_name}

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