jax-ml/jax · error · NotImplementedError

Unsupported layout for conversion from MLIR attribute: {attr

Error message

Unsupported layout for conversion from MLIR attribute: {attr}

What it means

layouts_lib.from_layout_attr converts an MLIR layout attribute into a Python layout object; only WGSPlacedFrag, WGStridedFrag, splat, and Tiled layout attrs (plus their fragmented variants) are recognized. Any other attribute type raises NotImplementedError.

Source

Thrown at jax/experimental/mosaic/gpu/layouts.py:133

      return _to_splat_fragmented_layout_attr(layout)
    case fa.WGStridedFragLayout():
      return _to_strided_fragmented_layout_attr(layout)
    case fa.TiledLayout():
      return _to_tiled_layout_attr(layout)
    case _:
      assert_never(layout)


def from_layout_attr(attr: ir.Attribute) -> fa.FragmentedLayout:
  """Constructs a layout from an MLIR attribute."""
  if isinstance(attr, mgpu.WGSplatFragLayoutAttr):
    return _from_splat_fragmented_layout_attr(attr)
  elif isinstance(attr, mgpu.WGStridedFragLayoutAttr):
    return _from_strided_fragmented_layout_attr(attr)
  elif isinstance(attr, mgpu.TiledLayoutAttr):
    return _from_tiled_layout_attr(attr)
  else:
    raise NotImplementedError(
        f"Unsupported layout for conversion from MLIR attribute: {attr}"
    )


def splat_is_compatible_with_tiled(
    l1: fa.WGSplatFragLayout, l2: fa.TiledLayout
) -> bool:
  # A splat layout is compatible with a tiled layout up to replication if each
  # dimension in the shape of the splat layout is divisible by the corresponding
  # dimension in the base tile shape.
  s1, s2 = l1.shape, l2.base_tile_shape
  return all(d1 % d2 == 0 for d1, d2 in zip(s1, s2))


def to_transform_attr(
    transform: launch_context.MemRefTransform | mgpu.SwizzlingMode,
) -> ir.Attribute:
  if isinstance(transform, launch_context.TileTransform):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use one of the supported layout attributes (TiledLayoutAttr, WGSPlacedFragLayoutAttr, etc.)
2. If defining a new layout, also extend from_layout_attr with a converter
3. Align jax and any custom MLIR pass versions

def is_supported_layout_attr(attr):
    return isinstance(attr, (mgpu.TiledLayoutAttr, mgpu.WGSPlacedFragLayoutAttr,
                             mgpu.WGStridedFragLayoutAttr, mgpu.WGSplatFragLayoutAttr))

When it happens

Trigger: A custom or newly introduced mgpu layout attribute reaches from_layout_attr — e.g. when serializing/deserializing Mosaic kernels or lowering custom ops with a novel layout attr.

Common situations: Version mismatch where MLIR produces a layout attr the Python side doesn't know, or third-party dialects reusing mgpu layout attributes.

Related errors

• Replicated dimensions are not supported
• Loading multiple row tiles
• {uninitialized_memory} + {str(dtype)}
• Batching over dynamic grid values is not supported yet.
• dma_start not implemented in LoJAX yet.

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