jax-ml/jax · error · ValueError

{op} has an unsupported layout: {out_layout_attr}

Error message

{op} has an unsupported layout: {out_layout_attr}

What it means

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.

Source

Thrown at jax/experimental/mosaic/gpu/dialect_lowering.py:455

  if isinstance(out_layout, fa.WGStridedFragLayout):
    # TODO(bchetioui): Process transforms.
    if transforms_attr is not None:
      swizzle = swizzle_from_transforms_attr(transforms_attr)
      transforms = memref_transforms_from_transforms_attr(transforms_attr)
      if swizzle != mgpu.SwizzlingMode.kNoSwizzle or transforms:
        raise NotImplementedError(
            "Transformed or swizzled strided loads are not supported"
        )

    fragmented_array = fa.FragmentedArray.load_strided(
        transformed_ref,
        is_signed=is_signed,
        vec_size=out_layout.vec_size,
    )
    return [_fragmented_array_to_ir(fragmented_array)]

  if not isinstance(out_layout, fa.TiledLayout):
    raise ValueError(f"{op} has an unsupported layout: {out_layout_attr}")

  optimized = op.optimized.value if op.optimized is not None else None
  if transformed_ref.type.memory_space is None:  # GMEM
    fragmented_array = fa.FragmentedArray.load_untiled(
        transformed_ref,
        layout=out_layout,
        is_signed=is_signed,
        optimized=bool(optimized),
    )
    return [_fragmented_array_to_ir(fragmented_array)]

  if transforms_attr is None:
    raise ValueError(f"Unsupported memory space: {orig_ref_ty.memory_space}")

  swizzle = swizzle_from_transforms_attr(transforms_attr)
  transforms = memref_transforms_from_transforms_attr(transforms_attr)
  if transforms:
    [tiling_transform] = transforms

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. 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
2. Check the op's layout attribute with layouts_lib.from_layout_attr to confirm what was inferred
3. If a register layout is desired, load untiled/with a TiledLayout then convert via layout_cast

Example fix

// before
value = load(ref, layout=wg_frag_layout)
// after
value = load(ref, layout=tiled_layout)
value = value.to_layout(wg_frag_layout)

from jax.experimental.mosaic.gpu import layouts as layouts_lib
from jax.experimental.mosaic.gpu import fragmented_array as fa
layout = layouts_lib.from_layout_attr(out_layout_attr)
assert isinstance(layout, fa.TiledLayout), f'load needs TiledLayout, got {type(layout)}'

def is_tiled_layout(attr) -> bool:
    return isinstance(layouts_lib.from_layout_attr(attr), fa.TiledLayout)

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• Expected TiledLayout, got {type(layout)}
• Output layout {out_layout} must match the accumulator layout
• Unsupported layout: {src.layout}
• Replicated dimensions are not supported
• Cannot assign layout to async load with gather indices since

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