{"record":{"id":"d68f4e27d6117820","repo":"jax-ml/jax","slug":"unsupported-transforms-b-scale-transforms","errorCode":null,"errorMessage":"Unsupported transforms: {b_scale_transforms}","messagePattern":"Unsupported transforms: (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":2820,"sourceCode":"    )\n    if a_scale_transforms:\n      raise NotImplementedError(\n          f\"Unsupported transforms: {a_scale_transforms}\"\n      )\n  if b_scale_ref is not None and b_scale_transforms_tree is not None:\n    assert isinstance(b_scale_ref_aval, state.AbstractRef)\n    b_scale_transforms = b_scale_transforms_tree.unflatten(\n        b_scale_transforms_leaves\n    )\n    b_scale_transform_avals = b_scale_transforms_tree.unflatten(\n        b_scale_transforms_leaves_avals\n    )\n    b_scale_ref, _, b_scale_transforms = lowering._handle_transforms(\n        ctx, b_scale_ref_aval, b_scale_ref, b_scale_transform_avals,\n        b_scale_transforms\n    )\n    if b_scale_transforms:\n      raise NotImplementedError(f\"Unsupported transforms: {b_scale_transforms}\")\n  if a_sparse_metadata_transforms_tree is not None:\n    a_sparse_metadata_transforms = a_sparse_metadata_transforms_tree.unflatten(\n        a_sparse_metadata_transforms_leaves\n    )\n    a_sparse_metadata_transform_avals = (\n        a_sparse_metadata_transforms_tree.unflatten(\n            a_sparse_metadata_transforms_leaves_avals\n        )\n    )\n    assert isinstance(a_sparse_metadata_ref_aval, state_types.AbstractRef)\n    a_sparse_metadata_ref, _, a_sparse_metadata_transforms = (\n        lowering._handle_transforms(  # pyrefly: ignore[bad-specialization]\n            ctx, a_sparse_metadata_ref_aval, a_sparse_metadata_ref,\n            a_sparse_metadata_transform_avals, a_sparse_metadata_transforms)\n    )\n    if a_sparse_metadata_transforms:\n      raise NotImplementedError(\n          f\"Unsupported transforms: {a_sparse_metadata_transforms}\"","sourceCodeStart":2802,"sourceCodeEnd":2838,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L2802-L2838","documentation":"Same restriction as the A-scale case but for the B operand's scale factor reference: after lowering._handle_transforms runs, any residual transforms on b_scale_ref are unsupported and raise NotImplementedError.","triggerScenarios":"Passing b_scale_ref with remaining transforms (transposes etc.) in b_scale_transforms_tree to tcgen05_mma.","commonSituations":"Block-scaled (fp8) kernels where the B-scale tensor is accessed through a transformed reference.","solutions":["Eliminate transforms on the b_scale reference; load the scale already in the required layout","Materialize any transpose of the scale tensor manually before the MMA"],"exampleFix":"// before\ntcgen05_mma(a, b, acc, b_scale=scale_ref, b_scale_transforms=transforms)\n// after\ntcgen05_mma(a, b, acc, b_scale=scale_ref)","handlingStrategy":"validation","validationCode":"assert not b_scale_transforms or all(t is None for t in b_scale_transforms)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Materialize scale layouts manually in SMEM","Keep scale refs transform-free"],"tags":["jax","pallas","tcgen05","fp8","scaling","not-implemented"],"backgroundTag":"unsupported-transform-on-operand","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}