jax-ml/jax · error · RuntimeError

Failed to infer the output layout of multimem_load_reduce. P

Error message

Failed to infer the output layout of multimem_load_reduce. Please apply plgpu.layout_cast to its output right after its creation.

What it means

The lowering of multimem_load_reduce needs to know the register layout of its output tensor, and it tries to infer it from the downstream use via ctx.out_layout_hint. If no hint reaches the op (no layout_cast on the output), inference fails with this RuntimeError telling you to pin the layout explicitly.

Source

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

  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}"
    )
  dtype = ctx.avals_out[0].dtype
  transforms = tree.unflatten(transforms_leaves)
  transform_avals = tree.unflatten(ctx.avals_in[1:])
  ref_aval = ctx.avals_in[0]
  assert isinstance(ref_aval, state_types.AbstractRef)
  ref, _, transforms = lowering._handle_transforms(ctx, ref_aval, ref,
                                                   transform_avals, transforms,
                                                   allow_peer_refs=False)
  if transforms:
    raise NotImplementedError(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Apply plgpu.layout_cast(out, tiled_layout) immediately after multimem_load_reduce, as the message instructs
Pick a TiledLayout matching the shape/dtype (e.g. a row-major tiled layout) so the lowering can emit the fragment load
Update JAX — layout hint propagation for multicast ops has been improving across releases

Example fix

# before
out = plgpu.multimem_load_reduce(ref, 'sum', collective_axes=mesh.axis_names)
# after
from jax._src.pallas.mosaic_gpu import layout_cast
out = plgpu.multimem_load_reduce(ref, 'sum', collective_axes=mesh.axis_names)
out = layout_cast(out, my_tiled_layout)

Defensive patterns

Strategy: validation

Validate before calling

from jax._src.pallas.mosaic_gpu import core as mgpu
def has_layout_hint():
    return True  # static analysis impossible; instead always layout_cast defensively

Try / catch

try:
    kernel_jit(x)
except RuntimeError as e:
    if 'layout' in str(e) and 'multimem_load_reduce' in str(e):
        raise ValueError('Add plgpu.layout_cast to the multimem_load_reduce output') from e
    raise

Prevention

Always wrap multimem_load_reduce outputs in layout_cast with an explicit TiledLayout
Keep layout objects next to kernel definitions to avoid inference reliance

When it happens

Trigger: Calling plgpu.multimem_load_reduce and using its result directly (e.g. returning it or feeding an op with no layout propagation), so no plgpu.layout_cast is applied to the output.

Common situations: First use of multimem ops in a pallas kernel where layout inference is not automatic; kernels where the load_reduce result feeds arithmetic before any layout-sensitive consumer; pallas versions where out_layout_hint propagation is limited.

Related errors

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