{"record":{"id":"d7fb71341fb3318d","repo":"jax-ml/jax","slug":"failed-to-infer-the-output-layout-of-multimem-load","errorCode":null,"errorMessage":"Failed to infer the output layout of multimem_load_reduce. Please apply plgpu.layout_cast to its output right after its creation.","messagePattern":"Failed to infer the output layout of multimem_load_reduce\\. Please apply plgpu\\.layout_cast to its output right after its creation\\.","errorType":"exception","errorClass":"RuntimeError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":5445,"sourceCode":"  assert isinstance(out_ref, state_types.AbstractRef)\n  return out_ref.inner_aval, {pallas_core.comms_effect}\n\n@lowering.register_lowering_rule(multimem_load_reduce_p, mgpu.LoweringSemantics.Lane)\ndef _multimem_load_reduce_lowering_rule(\n    ctx: lowering.LoweringRuleContext, ref, *transforms_leaves, tree, collective_axes, reduction_op,\n):\n  if (mesh_info := ctx.module_ctx.mesh_info) is None:\n    raise ValueError(\n        \"JAX device mesh is required by multimem_load_reduce, but not defined.\"\n    )\n  if set(collective_axes) != set(mesh_info.axis_names):\n    raise NotImplementedError(\n        \"Only collective_axes that include all JAX device mesh\"\n        f\" ({mesh_info.axis_names}) axes are supported, but got\"\n        f\" {collective_axes}\"\n    )\n  if (layout := ctx.out_layout_hint) is None:\n    raise RuntimeError(\n        \"Failed to infer the output layout of multimem_load_reduce. Please apply\"\n        \" plgpu.layout_cast to its output right after its creation.\"\n    )\n  if not isinstance(layout, (mgpu.TiledLayout, mgpu.WGStridedFragLayout)):\n    raise ValueError(\n        \"Only tiled and WG strided layouts are supported by\"\n        f\" multimem_load_reduce, but got {layout}\"\n    )\n  dtype = ctx.avals_out[0].dtype\n  transforms = tree.unflatten(transforms_leaves)\n  transform_avals = tree.unflatten(ctx.avals_in[1:])\n  ref_aval = ctx.avals_in[0]\n  assert isinstance(ref_aval, state_types.AbstractRef)\n  ref, _, transforms = lowering._handle_transforms(ctx, ref_aval, ref,\n                                                   transform_avals, transforms,\n                                                   allow_peer_refs=False)\n  if transforms:\n    raise NotImplementedError(","sourceCodeStart":5427,"sourceCodeEnd":5463,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L5427-L5463","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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"],"exampleFix":"# before\nout = plgpu.multimem_load_reduce(ref, 'sum', collective_axes=mesh.axis_names)\n# after\nfrom jax._src.pallas.mosaic_gpu import layout_cast\nout = plgpu.multimem_load_reduce(ref, 'sum', collective_axes=mesh.axis_names)\nout = layout_cast(out, my_tiled_layout)","handlingStrategy":"validation","validationCode":"from jax._src.pallas.mosaic_gpu import core as mgpu\ndef has_layout_hint():\n    return True  # static analysis impossible; instead always layout_cast defensively","typeGuard":null,"tryCatchPattern":"try:\n    kernel_jit(x)\nexcept RuntimeError as e:\n    if 'layout' in str(e) and 'multimem_load_reduce' in str(e):\n        raise ValueError('Add plgpu.layout_cast to the multimem_load_reduce output') from e\n    raise","preventionTips":["Always wrap multimem_load_reduce outputs in layout_cast with an explicit TiledLayout","Keep layout objects next to kernel definitions to avoid inference reliance"],"tags":["jax","pallas","layout","multimem","tpu"],"backgroundTag":"layout-inference-failed","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}