{"record":{"id":"6217369f69abd289","repo":"jax-ml/jax","slug":"make-sure-that-the-axis-name-passed-to-jax-lax-ppe","errorCode":null,"errorMessage":"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.","messagePattern":"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\\.","errorType":"exception","errorClass":"RuntimeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":1180,"sourceCode":"mlir.register_lowering(\n    pmin_p, partial(_all_reduce_lowering, lax.min_p, lax.reduce_min))\nbatching.fancy_primitive_batchers[pmin_p] = \\\n  partial(_batched_reduction_collective, pmin_p, lambda v, axis_size: v)\n\n\ndef _pcollectives_lowering_common(ctx, *, axis_name, perm, op_name):\n  replica_groups = _replica_groups(ctx.module_context.axis_context, axis_name, None)\n  group_size = len(replica_groups[0])\n  srcs, dsts = unzip2((src % group_size, dst % group_size) for src, dst in perm)\n  if not (len(srcs) == len(set(srcs)) and len(dsts) == len(set(dsts))):\n    msg = f\"{op_name} sources and destinations must be unique, got {{}}.\"\n    raise ValueError(msg.format(perm))\n\n  full_perm = np.zeros((len(replica_groups), len(perm), 2), np.int64)\n  for i, grp in enumerate(replica_groups):\n    sorted_grp = tuple(sorted(grp))\n    if config.raise_on_ppermute_sort_diff.value and sorted_grp != grp:\n      raise RuntimeError(\n          \"Make sure that the axis_name passed to jax.lax.ppermute is in the\"\n          \" same order as the axis_names declared on the mesh. If you want to\"\n          \" allow different order, you can disable the check via `with\"\n          \" jax.raise_on_ppermute_sort_diff(False):` context manager.\")\n    for j, (src, dst) in enumerate(perm):\n      full_perm[i, j, 0] = grp[src]\n      full_perm[i, j, 1] = grp[dst]\n  full_perm = full_perm.reshape((-1, 2))\n\n  axis_context = ctx.module_context.axis_context\n  if isinstance(axis_context, SPMDAxisContext) and axis_context.manual_axes:\n    other_args: dict[str, Any] = dict(\n        channel_handle=hlo.ChannelHandle.get(\n            mlir.COLLECTIVE_CHANNEL_ID, mlir.DEVICE_TO_DEVICE_TYPE\n        )\n    )\n  else:\n    other_args = {}","sourceCodeStart":1162,"sourceCodeEnd":1198,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L1162-L1198","documentation":"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.","triggerScenarios":"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).","commonSituations":"Multi-axis meshes like Mesh(axes=('i','j')) and calling ppermute over ('j','i'); refactoring mesh axis order without updating call sites.","solutions":["Reorder the axis_name argument to match the order declared on the Mesh","Suppress intentionally via with jax.raise_on_ppermute_sort_diff(False): if the reordering is intended","Update all ppermute call sites after changing mesh axis order"],"exampleFix":"// before\nmesh = jax.sharding.Mesh(devices, axis_names=('i','j'))\nlax.ppermute(x, ('j','i'), perm)\n// after\nlax.ppermute(x, ('i','j'), perm)  # match mesh order","handlingStrategy":"validation","validationCode":"def check_axis_order(mesh, axis_name):\n    names = axis_name if isinstance(axis_name, (list, tuple)) else (axis_name,)\n    idx = [mesh.axis_names.index(n) for n in names]\n    assert idx == sorted(idx), 'axis_name order must match mesh axis_names order'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Keep one canonical mesh definition module","Add a test asserting ppermute axis order matches the mesh"],"tags":["jax","ppermute","mesh","ordering"],"backgroundTag":"ordering-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}