jax-ml/jax · error · NotImplementedError

Only collective_axes that include all JAX device mesh axes

Error message

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}

What it means

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.

Source

Thrown at jax/_src/pallas/mosaic_gpu/primitives.py:450

          ctx,
          transform.device_id,
          transform.device_id_type,
          transform_aval.device_id,
      )
      peer_id = lowering._ensure_ir_value(peer_id, jnp.int32)
      continue
    elif isinstance(transform, gpu_core.MulticastRef):
      if not supports_multicast:
        raise ValueError(
            "Multicast refs are not supported by this primitive."
        )
      if (mesh_info := ctx.module_ctx.mesh_info) is None:
        raise ValueError(
            "JAX device mesh is required by multicast copies, but not defined."
            " Use jax.set_mesh."
        )
      if set(transform.collective_axes) != set(mesh_info.axis_names):
        raise NotImplementedError(
            "Only collective_axes that include all JAX device mesh  axes are"
            f" supported, but got {transform.collective_axes}. Make sure to"
            f" pass collective_axes={mesh_info.axis_names}"
        )
      peer_id = mgpu.GLOBAL_BROADCAST
      continue
    elif isinstance(transform, indexing.NDIndexer):
      indexers.append(transform)
    else:
      raise NotImplementedError(
          "Non-indexing transforms on GMEM refs are not implemented.")
  if indexers:
    indexer = lowering.merge_indexers(indexers)
    gmem_slice = lowering._ndindexer_indices(indexer, allow_arrays=True)
  else:
    gmem_slice = ()
  return dict(
      gmem_slice=gmem_slice,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. 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)
2. Ensure jax.set_mesh / mesh context manager is active so mesh_info is populated
3. 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

Example fix

# before
copy_gmem_to_smem(src, dst, collective_axes=('data',))
# after
with jax.set_mesh(mesh):  # mesh.axis_names == ('data', 'model')
  copy_gmem_to_smem(src, dst, collective_axes=mesh.axis_names)

from jax.sharding import Mesh
assert set(collective_axes) == set(mesh.axis_names), (
  f'collective_axes {collective_axes} must equal mesh axes {mesh.axis_names}')

When it happens

Trigger: 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.

Common situations: 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).

Related errors

• Can't make a multicast reference into a peer reference.
• Can't make a peer reference into a multicast reference.
• collective_axes is not supported in pallas_call. Use plgpu.k
• Expected exactly one collective axis, got {collective_axes=}
• Only the TMA implementation supports collective copies

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