jax-ml/jax · error · ValueError

JAX device mesh is required by multimem_load_reduce, but not

Error message

JAX device mesh is required by multimem_load_reduce, but not defined.

What it means

The multimem_load_reduce op requires a JAX device mesh to know which devices participate in the multicast load, but the kernel was lowered with no mesh context (ctx.module_ctx.mesh_info is None). This ValueError is thrown by the Lane-semantics lowering rule before any layout work happens.

Source

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

multimem_load_reduce_p = jax_core.Primitive("multimem_load_reduce")

@multimem_load_reduce_p.def_effectful_abstract_eval
def _multimem_load_reduce_abstract_eval(ref, *avals_flat, tree, collective_axes, reduction_op):
  del collective_axes, reduction_op
  _check_ref(ref, "ref", gpu_core.GMEM)
  out_ref = ref
  if tree is not None:
    transforms = jax.tree.unflatten(tree, avals_flat)
    out_ref = state.transform_type(transforms, ref)
  assert isinstance(out_ref, state_types.AbstractRef)
  return out_ref.inner_aval, {pallas_core.comms_effect}

@lowering.register_lowering_rule(multimem_load_reduce_p, mgpu.LoweringSemantics.Lane)
def _multimem_load_reduce_lowering_rule(
    ctx: lowering.LoweringRuleContext, ref, *transforms_leaves, tree, collective_axes, reduction_op,
):
  if (mesh_info := ctx.module_ctx.mesh_info) is None:
    raise ValueError(
        "JAX device mesh is required by multimem_load_reduce, 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 (layout := ctx.out_layout_hint) is None:
    raise RuntimeError(
        "Failed to infer the output layout of multimem_load_reduce. Please apply"
        " plgpu.layout_cast to its output right after its creation."
    )
  if not isinstance(layout, (mgpu.TiledLayout, mgpu.WGStridedFragLayout)):
    raise ValueError(
        "Only tiled and WG strided layouts are supported by"
        f" multimem_load_reduce, but got {layout}"
    )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Wrap the kernel launch in a mesh-defining context such as sharding_map or jax's mesh utilities so mesh_info is populated
Define the mesh with the same axes you pass as collective_axes
If you do not actually need cross-device multicast, replace multimem_load_reduce with a plain load followed by an explicit reduction

Example fix

# before
out = kernel_jit(x)  # no mesh
# after
with sharding_map(kernel_jit, mesh, in_axes=..., out_axes=...):
  out = ...  # mesh_info now available for multimem_load_reduce

Defensive patterns

Strategy: validation

Validate before calling

assert mesh is not None, 'multimem ops require an active Mesh/sharding_map context'

Try / catch

try:
    kernel_jit(x)
except ValueError as e:
    if 'device mesh is required' in str(e):
        with sharding_map(kernel_jit, mesh, ...):
            ...

Prevention

Always launch multimem kernels inside sharding_map or a Mesh context
Add a startup assert that the mesh is defined before jitting distributed kernels

When it happens

Trigger: Calling plgpu.multimem_load_reduce(...) in a pallas kernel that was compiled/launched outside of a Mesh/sharding_map context, so module_ctx.mesh_info is None.

Common situations: Prototyping a pallas kernel with jax.jit alone (no Mesh), then adding multimem ops; forgetting to wrap the kernel in sharding_map or map/with_mesh; running on a single device without a defined mesh.

Related errors

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