jax-ml/jax · error · NotImplementedError

Unsupported transforms for ACC: {acc_transforms}.

Error message

Unsupported transforms for ACC: {acc_transforms}.

What it means

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.

Source

Thrown at jax/_src/pallas/mosaic_gpu/primitives.py:2696

  (
      acc_transforms_leaves_avals,
      a_transforms_leaves_avals,
      b_transforms_leaves_avals,
      barrier_transforms_leaves_avals,
      a_scale_transforms_leaves_avals,
      b_scale_transforms_leaves_avals,
      a_sparse_metadata_transforms_leaves_avals,
  ) = transforms_avals_lists

  if acc_transforms_tree is not None:
    acc_transforms = acc_transforms_tree.unflatten(acc_transforms_leaves)
    acc_transform_avals = acc_transforms_tree.unflatten(acc_transforms_leaves_avals)
    acc, _, acc_transforms = lowering._handle_transforms(
        ctx, acc_aval, acc, acc_transform_avals, acc_transforms,
        handle_transposes=False
    )
    if acc_transforms:
      raise NotImplementedError(
          f"Unsupported transforms for ACC: {acc_transforms}."
      )

  if a_transforms_tree is not None:
    a_transforms = a_transforms_tree.unflatten(a_transforms_leaves)
    a_out_ty = state_types.transform_type(a_transforms, a_aval)
    assert isinstance(a_out_ty, state_types.AbstractRef)
    a_dtype = a_out_ty.dtype
    a_transform_avals = a_transforms_tree.unflatten(a_transforms_leaves_avals)
    a_ref, _, a_transforms = lowering._handle_transforms(
        ctx, a_aval, a_ref, a_transform_avals, a_transforms,
        handle_transposes=False, handle_reshapes=True)
    match a_transforms:
      case (
          gpu_core.UnswizzleRef(lhs_swizzle),
          gpu_core.UntilingTransform(lhs_tiling),
      ):
        lhs_transpose = False

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove all transforms from the accumulator ref; use a plain TMEM ref
2. Apply any layout adaptation to the values read/written, not to the acc ref itself
3. Refactor so transforms are only attached to A/B operands

Example fix

# before
acc_t = swizzle_ref(acc_tmem)  # TransformedRef
tcgen05.mma(a, b, acc_t, k_dim=k)
# after
tcgen05.mma(a, b, acc_tmem, k_dim=k)  # plain TMEM ref

assert not isinstance(acc, pallas_core.TransformedRef) or not acc.transforms

def acc_is_plain(ref):
    return not getattr(ref, 'transforms', None)

When it happens

Trigger: Passing a TransformedRef accumulator (e.g. swizzled, transposed, or disjoint-transformed) to tcgen05.mma; storing a transform on the accumulator ref.

Common situations: Applying a swizzle/transpose helper to all refs in a loop including the accumulator; reusing operand transform plumbing for the TMEM accumulator.

Related errors

• Unsupported transforms for LHS: {a_transforms}.
• Unsupported transforms for RHS: {b_transforms}.
• Sparse metadata format not implemented for {operand_dtype=}
• Unsupported TMEM ref {ref}.
• Unsupported transforms: {a_scale_transforms}

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/442fb59afd540d9d. Report an issue: GitHub.