jax-ml/jax · error · RuntimeError

Make sure that the axis_name passed to jax.lax.ppermute is i

Error message

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.

What it means

ppermute's lowering verifies that the device order of each replica group matches sorted order; a mismatch usually means the axis_name order in the call differs from the mesh's axis_names declaration. JAX raises (optionally, controlled by raise_on_ppermute_sort_diff) because XLA's source_target_pairs assume sorted group order.

Source

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

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(
        channel_handle=hlo.ChannelHandle.get(
            mlir.COLLECTIVE_CHANNEL_ID, mlir.DEVICE_TO_DEVICE_TYPE
        )
    )
  else:
    other_args = {}

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reorder the axis_name argument to match the order declared on the Mesh
2. Suppress intentionally via with jax.raise_on_ppermute_sort_diff(False): if the reordering is intended
3. Update all ppermute call sites after changing mesh axis order

Example fix

// before
mesh = jax.sharding.Mesh(devices, axis_names=('i','j'))
lax.ppermute(x, ('j','i'), perm)
// after
lax.ppermute(x, ('i','j'), perm)  # match mesh order

def check_axis_order(mesh, axis_name):
    names = axis_name if isinstance(axis_name, (list, tuple)) else (axis_name,)
    idx = [mesh.axis_names.index(n) for n in names]
    assert idx == sorted(idx), 'axis_name order must match mesh axis_names order'

When it happens

Trigger: Calling lax.ppermute with an axis_name tuple whose ordering differs from the mesh definition order, with the raise_on_ppermute_sort_diff config enabled (default).

Common situations: Multi-axis meshes like Mesh(axes=('i','j')) and calling ppermute over ('j','i'); refactoring mesh axis order without updating call sites.

Related errors

• Mapped away dimension of inputs passed to vmap should be sha
• Unmapped values passed to vmap cannot be sharded along the m
• callbacks are only supported in spmd computations when all m
• Mesh must be provided for shard_map with checkify.
• Sharding spec {spec} implies that array axis {dim} is partit

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