jax-ml/jax · error · NotImplementedError

Unsupported transform: {transform}

Error message

Unsupported transform: {transform}

What it means

memref_transforms_from_transforms_attr only understands TileTransformAttr and SwizzleTransformAttr; any other transform attribute is rejected with NotImplementedError.

Source

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

  swizzle = None
  for transform in attr:
    if isinstance(transform, mgpu.SwizzleTransformAttr):
      if swizzle is not None:
        raise ValueError("Found multiple SwizzleTransformAttr")
      swizzle = mgpu.SwizzlingMode(mgpu.SwizzleTransformAttr(transform).swizzle)
  return swizzle or mgpu.SwizzlingMode.kNoSwizzle


def memref_transforms_from_transforms_attr(
    attr: ir.ArrayAttr,
) -> tuple[lc.MemRefTransform, ...]:
  gmem_transforms: list[lc.MemRefTransform] = []
  for transform in attr:
    if isinstance(transform, mgpu.TileTransformAttr):
      tile_transform = lc.TileTransform(tuple(transform.tiling))
      gmem_transforms.append(tile_transform)
    elif not isinstance(transform, mgpu.SwizzleTransformAttr):
      raise NotImplementedError(f"Unsupported transform: {transform}")
  return tuple(gmem_transforms)


def tile_offset(
    offsets: tuple[int, ...], tiling: tuple[int, ...]
) -> tuple[int, ...]:
  """Tiles the trailing offsets in `offsets` according to `tiling`.

  Raises if the offsets are not aligned with the start of a tile.
  """
  if len(offsets) < len(tiling):
    raise ValueError(f"Offsets {offsets} have lower rank than tiling {tiling}")
  untiled_offsets, tiled_offsets = (
      offsets[: -len(tiling)],
      offsets[-len(tiling) :],
  )
  for i, t in zip(tiled_offsets, tiling, strict=True):
    if i % t != 0:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove the unsupported transform from the attribute
2. Upgrade JAX to a version that supports the transform kind
3. Express the memory layout with tile+swizzle only

for t in transforms:
    assert isinstance(t, (mgpu.TileTransformAttr, mgpu.SwizzleTransformAttr)), f'unsupported transform {t}'

When it happens

Trigger: Including a custom or newer transform attr (not tile/swizzle) in an op's transforms array.

Common situations: Version skew: newer MLIR/Mosaic dialects add transform kinds the installed JAX doesn't lower.

Understand the failure class

Background: UnsupportedOperationException and "is not supported" errors: when a library deliberately refuses a call — this error's family across 30 libraries.

Related errors

• Found multiple SwizzleTransformAttr
• Mosaic GPU does not yet support AMD ROCm devices. Use ``comp
• {op} has an unsupported layout: {out_layout_attr}
• Unsupported memory space: {orig_ref_ty.memory_space}
• Unsupported memory space: {ref_type.memory_space}

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