{"record":{"id":"f0c20051c339acc3","repo":"jax-ml/jax","slug":"only-collective-axes-that-include-all-jax-device-m-f0c200","errorCode":null,"errorMessage":"Only collective_axes that include all JAX device mesh ({mesh_info.axis_names}) axes are supported, but got {collective_axes}","messagePattern":"Only collective_axes that include all JAX device mesh \\((.+?)\\) axes are supported, but got (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":5378,"sourceCode":"    dtype = ty.dtype\n  if source.dtype != dtype:\n    raise ValueError(f\"Value dtype {source.dtype} does not match ref dtype {dtype}\")\n  if source.shape != shape:\n    raise ValueError(f\"Value shape {source.shape} does not match ref shape {shape}\")\n  return [], {pallas_core.comms_effect, state.WriteEffect(1)}\n\n\n@lowering.register_lowering_rule(multimem_store_p, mgpu.LoweringSemantics.Lane)\n@lowering.register_lowering_rule(multimem_store_p, mgpu.LoweringSemantics.Warpgroup)\ndef _multimem_store_lowering_rule(\n    ctx: lowering.LoweringRuleContext, value, local_ref, *transforms_leaves, transforms_tree, collective_axes,\n):\n  if (mesh_info := ctx.module_ctx.mesh_info) is None:\n    raise ValueError(\n        \"JAX device mesh is required by multimem_store, but not defined.\"\n    )\n  if set(collective_axes) != set(mesh_info.axis_names):\n    raise NotImplementedError(\n        \"Only collective_axes that include all JAX device mesh\"\n        f\" ({mesh_info.axis_names}) axes are supported, but got\"\n        f\" {collective_axes}\"\n    )\n  if transforms_tree is not None:\n    transforms = tree_util.tree_unflatten(transforms_tree, transforms_leaves)\n    local_ref_aval = ctx.avals_in[1]\n    assert isinstance(local_ref_aval, state_types.AbstractRef)\n    transform_avals = transforms_tree.unflatten(ctx.avals_in[2:])\n    local_ref, _, transforms = lowering._handle_transforms(\n        ctx, local_ref_aval, local_ref, transform_avals, transforms, allow_peer_refs=False\n    )\n    if transforms:\n      raise NotImplementedError(\n          f\"Unhandled transforms for multimem_store: {transforms}\"\n      )\n  multi_ref = ctx.launch_ctx.to_remote_multicast(local_ref)\n  scalar = not ctx.avals_in[0].shape","sourceCodeStart":5360,"sourceCodeEnd":5396,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L5360-L5396","documentation":"Raised during lowering of multimem_store on a TPU/GPU with a JAX device mesh. The multimem (multicast/reduced) hardware path only works when the collective_axes you pass cover every axis of the mesh the kernel was launched under. If the set of collective_axes differs from mesh_info.axis_names, this NotImplementedError fires.","triggerScenarios":"Calling plgpu.multimem_store(..., collective_axes=('data',)) inside a pallas kernel while jax.sharding.Mesh only defines axis 'data', or passing a subset/superset of the mesh axis names (e.g. mesh has ('data','model') but collective_axes=('data',)).","commonSituations":"Running distributed pallas kernels under jax.jit with sharding_map or a Mesh context; mismatch between the mesh you shard inputs with and the collective_axes passed to the multimem op; renaming mesh axes after refactoring a pipeline.","solutions":["Pass collective_axes equal to the full set of mesh axis names, e.g. collective_axes=tuple(mesh.axis_names), so both sets match","Run the kernel inside the same Mesh/sharding context that defines mesh_info (e.g. sharding_map or jax.lax.with_sharding_constraint setup) so mesh_info.axis_names matches your intended axes","If you only want a subset of devices to multicast, restructure the mesh so its axis names are exactly the axes you want to multicast over","Fall back to regular store + explicit collective (ppermute/all_reduce) if partial-axis multimem is required"],"exampleFix":"# before\nplgpu.multimem_store(ref, value, collective_axes=('data',))\n# after (mesh axis_names = ('data', 'model'))\nplgpu.multimem_store(ref, value, collective_axes=('data', 'model'))","handlingStrategy":"validation","validationCode":"import jax\nmesh = jax.create_mesh(...)  # your mesh\nassert set(collective_axes) == set(mesh.axis_names), (\n  f'collective_axes {collective_axes} must equal mesh axes {mesh.axis_names}')","typeGuard":"def axes_match(collective_axes: tuple[str, ...], mesh) -> bool:\n    return isinstance(collective_axes, tuple) and set(collective_axes) == set(mesh.axis_names)","tryCatchPattern":"try:\n    kernel_jit(x)\nexcept NotImplementedError as e:\n    if 'collective_axes' in str(e):\n        raise ValueError(f'Fix collective_axes: {e}') from e\n    raise","preventionTips":["Always derive collective_axes from mesh.axis_names instead of hardcoding","Keep one source of truth for mesh axis names in your distributed setup"],"tags":["jax","pallas","tpu","distributed","device-mesh","multimem"],"backgroundTag":"mesh-axis-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}