{"record":{"id":"dce72830187ba681","repo":"jax-ml/jax","slug":"only-tiled-and-wg-strided-layouts-are-supported-by","errorCode":null,"errorMessage":"Only tiled and WG strided layouts are supported by multimem_load_reduce, but got {layout}","messagePattern":"Only tiled and WG strided layouts are supported by multimem_load_reduce, but got (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":5450,"sourceCode":"    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(\n        f\"Unhandled transforms for multimem_load_reduce: {transforms}\"\n    )\n  multi_ref = ctx.launch_ctx.to_remote_multicast(ref)\n  is_signed = mgpu_utils.is_signed(dtype)\n  arr = mgpu.FragmentedArray.load_reduce_untiled(","sourceCodeStart":5432,"sourceCodeEnd":5468,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L5432-L5468","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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"],"exampleFix":"# before\nout = plgpu.layout_cast(out, strided_frag_layout)  # not supported\n# after\nout = plgpu.layout_cast(out, plgpu.TiledLayout(...))","handlingStrategy":"type-guard","validationCode":"from jax._src.pallas.mosaic_gpu import core as mgpu\nassert isinstance(layout, (mgpu.TiledLayout, mgpu.WGStridedFragLayout)), f'bad layout: {layout}'","typeGuard":"def is_supported_multimem_layout(layout) -> bool:\n    from jax._src.pallas.mosaic_gpu import core as mgpu\n    return isinstance(layout, (mgpu.TiledLayout, mgpu.WGStridedFragLayout))","tryCatchPattern":"try:\n    kernel_jit(x)\nexcept ValueError as e:\n    if 'tiled and WG strided' in str(e):\n        out = plgpu.layout_cast(out, plgpu.TiledLayout(...))\n    else:\n        raise","preventionTips":["Use a shared helper that only builds TiledLayout/WGStridedFragLayout for multimem outputs","Never reuse operand layouts from unrelated ops on multimem outputs"],"tags":["jax","pallas","layout","multimem","unsupported-layout"],"backgroundTag":"unsupported-register-layout","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}