{"record":{"id":"f557efe7e0fcab49","repo":"jax-ml/jax","slug":"unsupported-transforms-for-lhs-a-transforms","errorCode":null,"errorMessage":"Unsupported transforms for LHS: {a_transforms}.","messagePattern":"Unsupported transforms for LHS: (.+?)\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":2724,"sourceCode":"    a_ref, _, a_transforms = lowering._handle_transforms(\n        ctx, a_aval, a_ref, a_transform_avals, a_transforms,\n        handle_transposes=False, handle_reshapes=True)\n    match a_transforms:\n      case (\n          gpu_core.UnswizzleRef(lhs_swizzle),\n          gpu_core.UntilingTransform(lhs_tiling),\n      ):\n        lhs_transpose = False\n      case (\n          gpu_core.UnswizzleRef(lhs_swizzle),\n          gpu_core.UntilingTransform(lhs_tiling),\n          state_types.TransposeTransform((1, 0)),\n      ):\n        lhs_transpose = True\n      case () if isinstance(a_ref, tcgen05.TMEMRef):\n        lhs_tiling = None\n      case _:\n        raise NotImplementedError(\n            f\"Unsupported transforms for LHS: {a_transforms}.\"\n        )\n    if not isinstance(a_ref, tcgen05.TMEMRef):\n      assert lhs_swizzle is not None\n      swizzle_elems = 8 * lhs_swizzle // dtypes.itemsize_bits(a_dtype)\n      if lhs_tiling != (8, swizzle_elems):\n        raise ValueError(\"MMA lhs tiling does not fit swizzle. \"\n                        f\"{lhs_tiling=} expected={(8, swizzle_elems)}\")\n\n  assert b_transforms_tree is not None\n  b_transforms = b_transforms_tree.unflatten(b_transforms_leaves)\n  b_out_ty = state_types.transform_type(b_transforms, b_aval)\n  assert isinstance(b_out_ty, state_types.AbstractRef)\n  b_dtype = b_out_ty.dtype\n  b_transform_avals = b_transforms_tree.unflatten(b_transforms_leaves_avals)\n  b_ref, _, b_transforms = lowering._handle_transforms(\n      ctx, b_aval, b_ref, b_transform_avals, b_transforms, handle_transposes=False,\n      handle_reshapes=True)","sourceCodeStart":2706,"sourceCodeEnd":2742,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L2706-L2742","documentation":"At lowering, the LHS ref's transform stack must match one of the supported patterns: (unswizzle + untiling (+ transpose)) or empty for TMEM refs. Anything else raises NotImplementedError.","triggerScenarios":"Attaching an arbitrary/disjoint transform or a non-standard swizzle combination to the LHS ref of tcgen05.mma.","commonSituations":"Custom swizzle helpers composing transforms in a different order; using a DisjointTransform or sliced LHS ref.","solutions":["Use the standard swizzle helper (e.g. swizzle=128) that produces UnswizzleRef+UntilingTransform","Transpose via TransposeTransform((1,0)) if a transposed LHS is needed","For TMEM LHS, pass the ref with no transforms"],"exampleFix":"# before\na_weird = a_ref.transpose((1,0)).disjoint(...)  # unsupported stack\ntcgen05.mma(a_weird, b, acc, k_dim=k)\n# after\na_smem = load_to_smem(a, swizzle=128)\ntcgen05.mma(a_smem, b, acc, k_dim=k)","handlingStrategy":"validation","validationCode":"allowed = lambda t: t in (\n    (gpu_core.UnswizzleRef(swz), gpu_core.UntilingTransform(tl), state_types.TransposeTransform((1,0))),\n    (gpu_core.UnswizzleRef(swz), gpu_core.UntilingTransform(tl)),\n    ())\nassert a_transforms in allowed or isinstance(a.ref if hasattr(a,'ref') else a, tcgen05.TMEMRef)","typeGuard":"def lhs_transforms_supported(a_transforms):\n    for t in a_transforms:\n        if not isinstance(t, (gpu_core.UnswizzleRef, gpu_core.UntilingTransform, state_types.TransposeTransform)):\n            return False\n    return True","tryCatchPattern":null,"preventionTips":["Use standard swizzle/load utilities instead of manual transforms","Limit LHS transforms to unswizzle+untiling(+transpose)"],"tags":["jax","pallas","tcgen05","transform","not-implemented"],"backgroundTag":"unsupported-ref-transform","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}