jax-ml/jax · error · NotImplementedError

Only collective_axes that include all JAX device mesh ({mesh

Error message

Only collective_axes that include all JAX device mesh ({mesh_info.axis_names}) axes are supported, but got {collective_axes}

What it means

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.

Source

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

    dtype = ty.dtype
  if source.dtype != dtype:
    raise ValueError(f"Value dtype {source.dtype} does not match ref dtype {dtype}")
  if source.shape != shape:
    raise ValueError(f"Value shape {source.shape} does not match ref shape {shape}")
  return [], {pallas_core.comms_effect, state.WriteEffect(1)}


@lowering.register_lowering_rule(multimem_store_p, mgpu.LoweringSemantics.Lane)
@lowering.register_lowering_rule(multimem_store_p, mgpu.LoweringSemantics.Warpgroup)
def _multimem_store_lowering_rule(
    ctx: lowering.LoweringRuleContext, value, local_ref, *transforms_leaves, transforms_tree, collective_axes,
):
  if (mesh_info := ctx.module_ctx.mesh_info) is None:
    raise ValueError(
        "JAX device mesh is required by multimem_store, but not defined."
    )
  if set(collective_axes) != set(mesh_info.axis_names):
    raise NotImplementedError(
        "Only collective_axes that include all JAX device mesh"
        f" ({mesh_info.axis_names}) axes are supported, but got"
        f" {collective_axes}"
    )
  if transforms_tree is not None:
    transforms = tree_util.tree_unflatten(transforms_tree, transforms_leaves)
    local_ref_aval = ctx.avals_in[1]
    assert isinstance(local_ref_aval, state_types.AbstractRef)
    transform_avals = transforms_tree.unflatten(ctx.avals_in[2:])
    local_ref, _, transforms = lowering._handle_transforms(
        ctx, local_ref_aval, local_ref, transform_avals, transforms, allow_peer_refs=False
    )
    if transforms:
      raise NotImplementedError(
          f"Unhandled transforms for multimem_store: {transforms}"
      )
  multi_ref = ctx.launch_ctx.to_remote_multicast(local_ref)
  scalar = not ctx.avals_in[0].shape

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass collective_axes equal to the full set of mesh axis names, e.g. collective_axes=tuple(mesh.axis_names), so both sets match
2. 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
3. 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
4. Fall back to regular store + explicit collective (ppermute/all_reduce) if partial-axis multimem is required

Example fix

# before
plgpu.multimem_store(ref, value, collective_axes=('data',))
# after (mesh axis_names = ('data', 'model'))
plgpu.multimem_store(ref, value, collective_axes=('data', 'model'))

import jax
mesh = jax.create_mesh(...)  # your mesh
assert set(collective_axes) == set(mesh.axis_names), (
  f'collective_axes {collective_axes} must equal mesh axes {mesh.axis_names}')

def axes_match(collective_axes: tuple[str, ...], mesh) -> bool:
    return isinstance(collective_axes, tuple) and set(collective_axes) == set(mesh.axis_names)

try:
    kernel_jit(x)
except NotImplementedError as e:
    if 'collective_axes' in str(e):
        raise ValueError(f'Fix collective_axes: {e}') from e
    raise

When it happens

Trigger: 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',)).

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

Related errors

• JAX device mesh is required by multimem_load_reduce, but not
• Failed to infer the output layout of multimem_load_reduce. P
• Compiler params for platform {platform} cannot be used for {
• Memory space {self.memory_space} is not supported by mesh {s
• Acc ref must be at least 2D, got shape {shape}

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