{"record":{"id":"4521bcb0a3e2aaca","repo":"jax-ml/jax","slug":"unsupported-transforms-transforms-4521bc","errorCode":null,"errorMessage":"Unsupported transforms {transforms}","messagePattern":"Unsupported transforms (.+?)","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/layout_inference.py","lineNumber":2137,"sourceCode":") -> cs.SMEMTransforms:\n  transforms = [layouts_lib.from_transform_attr(x) for x in transform_attrs]\n  match transforms:\n    case []:\n      tile_transform = None\n      swizzle = None\n    case [lc.TileTransform() as t]:\n      tile_transform = t\n      swizzle = None\n    case [lc.TileTransform() as t, mgpu.SwizzlingMode() as s]:\n      tile_transform = t\n      # TODO(olechwierowicz): We should eliminate `kNoSwizzle`, representing\n      # this state as None is enough.\n      swizzle = s.value if s != mgpu.SwizzlingMode.kNoSwizzle else None\n    case [mgpu.SwizzlingMode() as s]:\n      tile_transform = None\n      swizzle = s.value if s != mgpu.SwizzlingMode.kNoSwizzle else None\n    case _:\n      raise NotImplementedError(f\"Unsupported transforms {transforms}\")\n\n  return cs.SMEMTransforms(tile_transform, swizzle)\n\n\n@_add_constraint_system_derivation_rule(mgpu.WithTransformsOp)\ndef _with_transforms_constraint_system(\n    ctx: DerivationContext,\n    op: mgpu.WithTransformsOp,\n) -> ConstraintSystemDerivationRuleResult:\n  source = ValueSite(op, VariableType.OPERAND, 0)\n  dest = ValueSite(op, VariableType.RESULT, 0)\n  var = ctx.producer_ref(source)\n  smem_transforms = _extract_smem_transforms_from_custom_transform_attrs(\n      op.transforms\n  )\n\n  if not cs.is_valid_assignment(var, smem_transforms):\n    tiling_transform_str = (","sourceCodeStart":2119,"sourceCodeEnd":2155,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/layout_inference.py#L2119-L2155","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Inspect the printed transforms list and reduce it to at most one tiling vector and one SwizzlingMode (drop SwizzlingMode.kNoSwizzle)","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","Upgrade jax/mosaic if you are passing a transform kind added in a newer version"],"exampleFix":"# before\nop = mgpu.with_transforms(src, transforms=[sw1, sw2])  # two swizzles -> unsupported\n\n# after\nop = mgpu.with_transforms(src, transforms=[sw1])  # at most one tiling + one swizzle","handlingStrategy":"validation","validationCode":"# transforms must be (), ([tiling...], swizzle_mode), or (swizzle_mode,)\nassert len(transforms) <= 2, 'too many transforms'\nassert sum(isinstance(t, mgpu.SwizzlingMode) for t in transforms) <= 1","typeGuard":"def is_supported_transforms(transforms) -> bool:\n    sw = [t for t in transforms if isinstance(t, mgpu.SwizzlingMode)]\n    til = [t for t in transforms if not isinstance(t, mgpu.SwizzlingMode)]\n    return len(sw) <= 1 and len(til) <= 1 and len(transforms) <= 2","tryCatchPattern":null,"preventionTips":["Build transform attrs via Mosaic helpers, never hand-assembled attribute lists","Keep transform format ([tile], swizzle) or ([], swizzle); drop kNoSwizzle"],"tags":["jax","mosaic-gpu","smem-transforms","swizzle","tiling","layout-inference"],"backgroundTag":"unsupported-attribute-value","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}