jax-ml/jax · error · NotImplementedError

Unsupported transforms {transforms}

Error message

Unsupported transforms {transforms}

What it means

Mosaic GPU's layout inference parses the custom transform attributes (swizzle/tiling) attached to ops like with_transforms or custom primitives. The match on the transform list handles only recognized patterns (empty, a tiling transform, a swizzle mode, or tiling+swizzle); any other combination raises NotImplementedError with the raw transforms printed.

Source

Thrown at jax/experimental/mosaic/gpu/layout_inference.py:2137

) -> cs.SMEMTransforms:
  transforms = [layouts_lib.from_transform_attr(x) for x in transform_attrs]
  match transforms:
    case []:
      tile_transform = None
      swizzle = None
    case [lc.TileTransform() as t]:
      tile_transform = t
      swizzle = None
    case [lc.TileTransform() as t, mgpu.SwizzlingMode() as s]:
      tile_transform = t
      # TODO(olechwierowicz): We should eliminate `kNoSwizzle`, representing
      # this state as None is enough.
      swizzle = s.value if s != mgpu.SwizzlingMode.kNoSwizzle else None
    case [mgpu.SwizzlingMode() as s]:
      tile_transform = None
      swizzle = s.value if s != mgpu.SwizzlingMode.kNoSwizzle else None
    case _:
      raise NotImplementedError(f"Unsupported transforms {transforms}")

  return cs.SMEMTransforms(tile_transform, swizzle)


@_add_constraint_system_derivation_rule(mgpu.WithTransformsOp)
def _with_transforms_constraint_system(
    ctx: DerivationContext,
    op: mgpu.WithTransformsOp,
) -> ConstraintSystemDerivationRuleResult:
  source = ValueSite(op, VariableType.OPERAND, 0)
  dest = ValueSite(op, VariableType.RESULT, 0)
  var = ctx.producer_ref(source)
  smem_transforms = _extract_smem_transforms_from_custom_transform_attrs(
      op.transforms
  )

  if not cs.is_valid_assignment(var, smem_transforms):
    tiling_transform_str = (

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Inspect the printed transforms list and reduce it to at most one tiling vector and one SwizzlingMode (drop SwizzlingMode.kNoSwizzle)
2. Construct transforms through Mosaic's public helpers (e.g. mgpu.smem_transforms / the transforms tuple format: ([...tile...], swizzle) or ([], swizzle)) rather than hand-building attributes
3. Upgrade jax/mosaic if you are passing a transform kind added in a newer version

Example fix

# before
op = mgpu.with_transforms(src, transforms=[sw1, sw2])  # two swizzles -> unsupported

# after
op = mgpu.with_transforms(src, transforms=[sw1])  # at most one tiling + one swizzle

# transforms must be (), ([tiling...], swizzle_mode), or (swizzle_mode,)
assert len(transforms) <= 2, 'too many transforms'
assert sum(isinstance(t, mgpu.SwizzlingMode) for t in transforms) <= 1

def is_supported_transforms(transforms) -> bool:
    sw = [t for t in transforms if isinstance(t, mgpu.SwizzlingMode)]
    til = [t for t in transforms if not isinstance(t, mgpu.SwizzlingMode)]
    return len(sw) <= 1 and len(til) <= 1 and len(transforms) <= 2

When it happens

Trigger: Attaching a transform attribute list containing entries other than a single MemrefTranspose-free tiling vector and/or a single SwizzlingMode — e.g. multiple swizzles, unknown attribute types, or reordered/malformed lists — to an op handled by _custom_primitive_constraint_system or _with_transforms_constraint_system.

Common situations: Hand-written MLIR or low-level Mosaic code constructing with_transforms with invalid transform lists; version mismatches where a newer/older transform attribute kind is not recognized by this Mosaic version.

Related errors

• Cannot apply {tiling_transform_str} with {swizzle_str} to me
• Transpose cannot be moved before a tiling transform when it
• Commuting a `UntilingTransform` with a `ReshapeTransform` is
• Folding tiled dimensions into untiled dimensions is not supp
• Cannot commute `UntilingTransform` with `ReshapeTransform` w

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