jax-ml/jax · error · ValueError

Only tiled and WG strided layouts are supported by multimem_

Error message

Only tiled and WG strided layouts are supported by multimem_load_reduce, but got {layout}

What it means

Even when an output layout hint exists, multimem_load_reduce can only materialize results in register layouts the multicast hardware supports: TiledLayout or WGStridedFragLayout. Any other layout (e.g. a generic strided fragment layout from another op) triggers this ValueError.

Source

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

    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(
        f"Unhandled transforms for multimem_load_reduce: {transforms}"
    )
  multi_ref = ctx.launch_ctx.to_remote_multicast(ref)
  is_signed = mgpu_utils.is_signed(dtype)
  arr = mgpu.FragmentedArray.load_reduce_untiled(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Cast the output to a TiledLayout (most common) built for the output shape/dtype
If on Warpgroup semantics and using fragment layouts, use WGStridedFragLayout explicitly
Do not propagate foreign layouts into the multimem_load_reduce output; insert layout_cast at the consumer instead

Example fix

# before
out = plgpu.layout_cast(out, strided_frag_layout)  # not supported
# after
out = plgpu.layout_cast(out, plgpu.TiledLayout(...))

Defensive patterns

Strategy: type-guard

Validate before calling

from jax._src.pallas.mosaic_gpu import core as mgpu
assert isinstance(layout, (mgpu.TiledLayout, mgpu.WGStridedFragLayout)), f'bad layout: {layout}'

Type guard

def is_supported_multimem_layout(layout) -> bool:
    from jax._src.pallas.mosaic_gpu import core as mgpu
    return isinstance(layout, (mgpu.TiledLayout, mgpu.WGStridedFragLayout))

Try / catch

try:
    kernel_jit(x)
except ValueError as e:
    if 'tiled and WG strided' in str(e):
        out = plgpu.layout_cast(out, plgpu.TiledLayout(...))
    else:
        raise

Prevention

Use a shared helper that only builds TiledLayout/WGStridedFragLayout for multimem outputs
Never reuse operand layouts from unrelated ops on multimem outputs

When it happens

Trigger: Applying plgpu.layout_cast to the output of multimem_load_reduce with a layout that is not mgpu.TiledLayout or mgpu.WGStridedFragLayout, and the hint propagates back to the lowering rule.

Common situations: Reusing a layout object produced for a different op (e.g. a plain strided fragment layout); layout_casting to the layout of an operand computed by dot or other non-tiled ops; mixing warpgroup and lane layout conventions.

Related errors

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