jax-ml/jax · error · NotImplementedError

Unsupported transform {transform}

Error message

Unsupported transform {transform}

What it means

to_transform_attr serializes a Python-side memory transform into an MLIR attribute, supporting only launch_context.TileTransform and mgpu.SwizzlingMode. Passing any other transform object raises NotImplementedError.

Source

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

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):
    return mgpu.TileTransformAttr.get(transform.tiling)
  elif isinstance(transform, mgpu.SwizzlingMode):
    return mgpu.SwizzleTransformAttr.get(transform)
  else:
    raise NotImplementedError(f"Unsupported transform {transform}")


def from_transform_attr(
    transform: ir.Attribute,
) -> launch_context.MemRefTransform | mgpu.SwizzlingMode:
  if isinstance(transform, mgpu.TileTransformAttr):
    return launch_context.TileTransform(
        tuple(mgpu.TileTransformAttr(transform).tiling)
    )
  elif isinstance(transform, mgpu.SwizzleTransformAttr):
    return mgpu.SwizzlingMode(mgpu.SwizzleTransformAttr(transform).swizzle)
  else:
    raise NotImplementedError(f"Unsupported transform {transform}")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Convert your transform to launch_context.TileTransform or mgpu.SwizzlingMode before building the layout
2. For new transform kinds, add a branch to to_transform_attr in a fork/PR

Example fix

// before
transforms=(my_custom_transform,)
// after
transforms=(launch_context.TileTransform((16, 16)), mgpu.SwizzlingMode.SWIZZLE_128B)

def is_supported_transform(t):
    return isinstance(t, (launch_context.TileTransform, mgpu.SwizzlingMode))

When it happens

Trigger: Calling layout-to-MLIR conversion (e.g. when emitting tiled layouts with transforms) with a transform object that is neither TileTransform nor SwizzlingMode — often a custom transform class.

Common situations: Extending Mosaic with custom swizzle/transform schemes, or passing a raw mgpu attr where a Python object is expected.

Related errors

• {uninitialized_memory} + {str(dtype)}
• Batching over dynamic grid values is not supported yet.
• dma_start not implemented in LoJAX yet.
• DMA partial discharge add=True not yet implemented.
• Unsupported transforms for ACC: {acc_transforms}.

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