{"record":{"id":"19962d3486944126","repo":"jax-ml/jax","slug":"op-has-an-unsupported-layout-out-layout-attr","errorCode":null,"errorMessage":"{op} has an unsupported layout: {out_layout_attr}","messagePattern":"(.+?) has an unsupported layout: (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/dialect_lowering.py","lineNumber":455,"sourceCode":"  if isinstance(out_layout, fa.WGStridedFragLayout):\n    # TODO(bchetioui): Process transforms.\n    if transforms_attr is not None:\n      swizzle = swizzle_from_transforms_attr(transforms_attr)\n      transforms = memref_transforms_from_transforms_attr(transforms_attr)\n      if swizzle != mgpu.SwizzlingMode.kNoSwizzle or transforms:\n        raise NotImplementedError(\n            \"Transformed or swizzled strided loads are not supported\"\n        )\n\n    fragmented_array = fa.FragmentedArray.load_strided(\n        transformed_ref,\n        is_signed=is_signed,\n        vec_size=out_layout.vec_size,\n    )\n    return [_fragmented_array_to_ir(fragmented_array)]\n\n  if not isinstance(out_layout, fa.TiledLayout):\n    raise ValueError(f\"{op} has an unsupported layout: {out_layout_attr}\")\n\n  optimized = op.optimized.value if op.optimized is not None else None\n  if transformed_ref.type.memory_space is None:  # GMEM\n    fragmented_array = fa.FragmentedArray.load_untiled(\n        transformed_ref,\n        layout=out_layout,\n        is_signed=is_signed,\n        optimized=bool(optimized),\n    )\n    return [_fragmented_array_to_ir(fragmented_array)]\n\n  if transforms_attr is None:\n    raise ValueError(f\"Unsupported memory space: {orig_ref_ty.memory_space}\")\n\n  swizzle = swizzle_from_transforms_attr(transforms_attr)\n  transforms = memref_transforms_from_transforms_attr(transforms_attr)\n  if transforms:\n    [tiling_transform] = transforms","sourceCodeStart":437,"sourceCodeEnd":473,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/dialect_lowering.py#L437-L473","documentation":"Raised by Mosaic GPU's lowering of mgpu.layout_cast/vector load ops when the output layout attribute on the operation is not a TiledLayout (e.g. a fragment/splat layout) where a tiled layout is required. The lowering path only supports loading into TiledLayouts from GMEM/SMEM.","triggerScenarios":"Calling a mosaic load (e.g. mosaic load with layout=) where the inferred out_layout attribute decodes to a non-TiledLayout such as WGStridedFragLayout or a Nanobind layout from layouts_lib.from_layout_attr.","commonSituations":"Passing a fragment layout (from a matmul accumulator) directly to a memory load; mixing tiled and register-fragment layouts in a kernel; upgrading Mosaic where layout inference changed.","solutions":["Use a tiled layout (e.g. layouts.TiledLayout or mosaic.tile) for the load's layout and cast afterwards if a fragment layout is needed","Check the op's layout attribute with layouts_lib.from_layout_attr to confirm what was inferred","If a register layout is desired, load untiled/with a TiledLayout then convert via layout_cast"],"exampleFix":"// before\nvalue = load(ref, layout=wg_frag_layout)\n// after\nvalue = load(ref, layout=tiled_layout)\nvalue = value.to_layout(wg_frag_layout)","handlingStrategy":"type-guard","validationCode":"from jax.experimental.mosaic.gpu import layouts as layouts_lib\nfrom jax.experimental.mosaic.gpu import fragmented_array as fa\nlayout = layouts_lib.from_layout_attr(out_layout_attr)\nassert isinstance(layout, fa.TiledLayout), f'load needs TiledLayout, got {type(layout)}'","typeGuard":"def is_tiled_layout(attr) -> bool:\n    return isinstance(layouts_lib.from_layout_attr(attr), fa.TiledLayout)","tryCatchPattern":null,"preventionTips":["Always load from memory with a TiledLayout; cast to fragment layouts afterwards","Keep a helper that asserts layout kinds before emitting mgpu ops"],"tags":["jax","mosaic","gpu","layout","tiled-layout"],"backgroundTag":"unsupported-layout","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}