{"record":{"id":"442fb59afd540d9d","repo":"jax-ml/jax","slug":"unsupported-transforms-for-acc-acc-transforms","errorCode":null,"errorMessage":"Unsupported transforms for ACC: {acc_transforms}.","messagePattern":"Unsupported transforms for ACC: (.+?)\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":2696,"sourceCode":"  (\n      acc_transforms_leaves_avals,\n      a_transforms_leaves_avals,\n      b_transforms_leaves_avals,\n      barrier_transforms_leaves_avals,\n      a_scale_transforms_leaves_avals,\n      b_scale_transforms_leaves_avals,\n      a_sparse_metadata_transforms_leaves_avals,\n  ) = transforms_avals_lists\n\n  if acc_transforms_tree is not None:\n    acc_transforms = acc_transforms_tree.unflatten(acc_transforms_leaves)\n    acc_transform_avals = acc_transforms_tree.unflatten(acc_transforms_leaves_avals)\n    acc, _, acc_transforms = lowering._handle_transforms(\n        ctx, acc_aval, acc, acc_transform_avals, acc_transforms,\n        handle_transposes=False\n    )\n    if acc_transforms:\n      raise NotImplementedError(\n          f\"Unsupported transforms for ACC: {acc_transforms}.\"\n      )\n\n  if a_transforms_tree is not None:\n    a_transforms = a_transforms_tree.unflatten(a_transforms_leaves)\n    a_out_ty = state_types.transform_type(a_transforms, a_aval)\n    assert isinstance(a_out_ty, state_types.AbstractRef)\n    a_dtype = a_out_ty.dtype\n    a_transform_avals = a_transforms_tree.unflatten(a_transforms_leaves_avals)\n    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","sourceCodeStart":2678,"sourceCodeEnd":2714,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L2678-L2714","documentation":"At lowering time, tcgen05.mma cannot apply any remaining transforms to the accumulator ref after transform handling (transposes are not handled for ACC). Any leftover transform raises NotImplementedError.","triggerScenarios":"Passing a TransformedRef accumulator (e.g. swizzled, transposed, or disjoint-transformed) to tcgen05.mma; storing a transform on the accumulator ref.","commonSituations":"Applying a swizzle/transpose helper to all refs in a loop including the accumulator; reusing operand transform plumbing for the TMEM accumulator.","solutions":["Remove all transforms from the accumulator ref; use a plain TMEM ref","Apply any layout adaptation to the values read/written, not to the acc ref itself","Refactor so transforms are only attached to A/B operands"],"exampleFix":"# before\nacc_t = swizzle_ref(acc_tmem)  # TransformedRef\ntcgen05.mma(a, b, acc_t, k_dim=k)\n# after\ntcgen05.mma(a, b, acc_tmem, k_dim=k)  # plain TMEM ref","handlingStrategy":"validation","validationCode":"assert not isinstance(acc, pallas_core.TransformedRef) or not acc.transforms","typeGuard":"def acc_is_plain(ref):\n    return not getattr(ref, 'transforms', None)","tryCatchPattern":null,"preventionTips":["Never attach swizzle/transpose helpers to the accumulator","Pass acc.ref when unwrapping a TransformedRef accumulator"],"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"}