{"record":{"id":"c53757dafdb30283","repo":"jax-ml/jax","slug":"unsupported-transforms-for-rhs-b-transforms","errorCode":null,"errorMessage":"Unsupported transforms for RHS: {b_transforms}.","messagePattern":"Unsupported transforms for RHS: (.+?)\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":2756,"sourceCode":"  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)\n  match b_transforms:\n    case (\n        gpu_core.UnswizzleRef(rhs_swizzle),\n        gpu_core.UntilingTransform(rhs_tiling),\n    ):\n      rhs_transpose = False\n    case (\n        gpu_core.UnswizzleRef(rhs_swizzle),\n        gpu_core.UntilingTransform(rhs_tiling),\n        state_types.TransposeTransform((1, 0)),\n    ):\n      rhs_transpose = True\n    case _:\n      raise NotImplementedError(\n          f\"Unsupported transforms for RHS: {b_transforms}.\"\n      )\n  swizzle_elems = 8 * rhs_swizzle // dtypes.itemsize_bits(b_dtype)\n  if rhs_tiling != (8, swizzle_elems):\n    raise ValueError(\n        \"MMA rhs tiling does not fit swizzle\"\n        f\" {rhs_tiling=} expected={(8, swizzle_elems)}\"\n    )\n\n  if barrier_transforms_tree is not None and barrier_ref is not None:\n    barrier_transforms = barrier_transforms_tree.unflatten(\n        barrier_transforms_leaves\n    )\n    base_index = _get_barrier_base_index(barrier_ref_aval, barrier_transforms)\n    if base_index is not None:\n      barrier_ref = barrier_ref[base_index]\n\n  if lhs_swizzle is None:","sourceCodeStart":2738,"sourceCodeEnd":2774,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L2738-L2774","documentation":"The RHS ref's transform stack in tcgen05.mma must match the supported pattern (unswizzle + untiling, optionally with transpose). Any other combination of transforms raises NotImplementedError.","triggerScenarios":"Attaching unsupported transforms (disjoint slicing, arbitrary transposes/permutations) to the B operand ref passed to tcgen05.mma.","commonSituations":"Slicing a larger SMEM buffer for B and passing the sliced TransformedRef; composing transforms in an order the lowering does not recognize.","solutions":["Load B into a dedicated SMEM buffer with the standard swizzle/tiling helper","Express transpose only via TransposeTransform((1,0))","Avoid slicing/disjoint transforms on the RHS ref; materialize the tile"],"exampleFix":"# before\nb_ref = big_smem[...]  # disjoint/sliced TransformedRef\ntcgen05.mma(a, b_ref, acc, k_dim=k)\n# after\nb_smem = load_to_smem(b, swizzle=128)\ntcgen05.mma(a, b_smem, acc, k_dim=k)","handlingStrategy":"validation","validationCode":"for t in b_transforms:\n    assert isinstance(t, (gpu_core.UnswizzleRef, gpu_core.UntilingTransform, state_types.TransposeTransform)), t","typeGuard":"def rhs_transforms_supported(b_transforms):\n    return all(isinstance(t, (gpu_core.UnswizzleRef, gpu_core.UntilingTransform, state_types.TransposeTransform)) for t in b_transforms)","tryCatchPattern":null,"preventionTips":["Materialize B tiles into dedicated SMEM buffers rather than slicing","Use the provided swizzle helpers"],"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"}