{"record":{"id":"782f063d3da7a8ed","repo":"jax-ml/jax","slug":"only-collective-axes-that-include-all-jax-device-m","errorCode":null,"errorMessage":"Only collective_axes that include all JAX device mesh  axes are supported, but got {transform.collective_axes}. Make sure to pass collective_axes={mesh_info.axis_names}","messagePattern":"Only collective_axes that include all JAX device mesh  axes are supported, but got (.+?)\\. Make sure to pass collective_axes=(.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":450,"sourceCode":"          ctx,\n          transform.device_id,\n          transform.device_id_type,\n          transform_aval.device_id,\n      )\n      peer_id = lowering._ensure_ir_value(peer_id, jnp.int32)\n      continue\n    elif isinstance(transform, gpu_core.MulticastRef):\n      if not supports_multicast:\n        raise ValueError(\n            \"Multicast refs are not supported by this primitive.\"\n        )\n      if (mesh_info := ctx.module_ctx.mesh_info) is None:\n        raise ValueError(\n            \"JAX device mesh is required by multicast copies, but not defined.\"\n            \" Use jax.set_mesh.\"\n        )\n      if set(transform.collective_axes) != set(mesh_info.axis_names):\n        raise NotImplementedError(\n            \"Only collective_axes that include all JAX device mesh  axes are\"\n            f\" supported, but got {transform.collective_axes}. Make sure to\"\n            f\" pass collective_axes={mesh_info.axis_names}\"\n        )\n      peer_id = mgpu.GLOBAL_BROADCAST\n      continue\n    elif isinstance(transform, indexing.NDIndexer):\n      indexers.append(transform)\n    else:\n      raise NotImplementedError(\n          \"Non-indexing transforms on GMEM refs are not implemented.\")\n  if indexers:\n    indexer = lowering.merge_indexers(indexers)\n    gmem_slice = lowering._ndindexer_indices(indexer, allow_arrays=True)\n  else:\n    gmem_slice = ()\n  return dict(\n      gmem_slice=gmem_slice,","sourceCodeStart":432,"sourceCodeEnd":468,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L432-L468","documentation":"Raised by _extract_gmem_copy_params when a multicast copy (collective_axes set) is requested but the set of collective axes doesn't exactly match all axes of the JAX device mesh (mesh_info.axis_names). The Mosaic GPU multicast path only supports broadcasting over the entire mesh; partial-axis multicast is not implemented.","triggerScenarios":"Calling copy_gmem_to_smem / copy_smem_to_gmem (or async prefetch) with collective_axes that is a subset, superset, or differently-named set than the mesh axes defined via jax.set_mesh; or using a Mesh whose axis_names don't match the collective_axes tuple passed to the Pallas kernel.","commonSituations":"Defining a mesh with multiple axes (e.g. ('data','model')) but passing collective_axes=('data',); renaming mesh axes between config and kernel code; forgetting to call jax.set_mesh at all (gives the sibling 'mesh required' error).","solutions":["Set collective_axes to exactly the full tuple of mesh axis names, e.g. collective_axes=mesh.axis_names (or read it programmatically inside the kernel)","Ensure jax.set_mesh / mesh context manager is active so mesh_info is populated","If you only want a subset of devices, reshape the Mesh so it has a single axis covering exactly the devices you want to multicast over"],"exampleFix":"# before\ncopy_gmem_to_smem(src, dst, collective_axes=('data',))\n# after\nwith jax.set_mesh(mesh):  # mesh.axis_names == ('data', 'model')\n  copy_gmem_to_smem(src, dst, collective_axes=mesh.axis_names)","handlingStrategy":"validation","validationCode":"from jax.sharding import Mesh\nassert set(collective_axes) == set(mesh.axis_names), (\n  f'collective_axes {collective_axes} must equal mesh axes {mesh.axis_names}')","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Derive collective_axes from mesh.axis_names instead of hardcoding","Always establish the mesh via jax.set_mesh before launching multicast kernels"],"tags":["jax","pallas","mosaic-gpu","collective-axes","device-mesh","multicast"],"backgroundTag":"axis-name-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}