jax-ml/jax · error · NotImplementedError

Unsupported transforms for RHS: {b_transforms}.

Error message

Unsupported transforms for RHS: {b_transforms}.

What it means

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.

Source

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

  b_dtype = b_out_ty.dtype
  b_transform_avals = b_transforms_tree.unflatten(b_transforms_leaves_avals)
  b_ref, _, b_transforms = lowering._handle_transforms(
      ctx, b_aval, b_ref, b_transform_avals, b_transforms, handle_transposes=False,
      handle_reshapes=True)
  match b_transforms:
    case (
        gpu_core.UnswizzleRef(rhs_swizzle),
        gpu_core.UntilingTransform(rhs_tiling),
    ):
      rhs_transpose = False
    case (
        gpu_core.UnswizzleRef(rhs_swizzle),
        gpu_core.UntilingTransform(rhs_tiling),
        state_types.TransposeTransform((1, 0)),
    ):
      rhs_transpose = True
    case _:
      raise NotImplementedError(
          f"Unsupported transforms for RHS: {b_transforms}."
      )
  swizzle_elems = 8 * rhs_swizzle // dtypes.itemsize_bits(b_dtype)
  if rhs_tiling != (8, swizzle_elems):
    raise ValueError(
        "MMA rhs tiling does not fit swizzle"
        f" {rhs_tiling=} expected={(8, swizzle_elems)}"
    )

  if barrier_transforms_tree is not None and barrier_ref is not None:
    barrier_transforms = barrier_transforms_tree.unflatten(
        barrier_transforms_leaves
    )
    base_index = _get_barrier_base_index(barrier_ref_aval, barrier_transforms)
    if base_index is not None:
      barrier_ref = barrier_ref[base_index]

  if lhs_swizzle is None:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Load B into a dedicated SMEM buffer with the standard swizzle/tiling helper
2. Express transpose only via TransposeTransform((1,0))
3. Avoid slicing/disjoint transforms on the RHS ref; materialize the tile

Example fix

# before
b_ref = big_smem[...]  # disjoint/sliced TransformedRef
tcgen05.mma(a, b_ref, acc, k_dim=k)
# after
b_smem = load_to_smem(b, swizzle=128)
tcgen05.mma(a, b_smem, acc, k_dim=k)

for t in b_transforms:
    assert isinstance(t, (gpu_core.UnswizzleRef, gpu_core.UntilingTransform, state_types.TransposeTransform)), t

def rhs_transforms_supported(b_transforms):
    return all(isinstance(t, (gpu_core.UnswizzleRef, gpu_core.UntilingTransform, state_types.TransposeTransform)) for t in b_transforms)

When it happens

Trigger: Attaching unsupported transforms (disjoint slicing, arbitrary transposes/permutations) to the B operand ref passed to tcgen05.mma.

Common situations: Slicing a larger SMEM buffer for B and passing the sliced TransformedRef; composing transforms in an order the lowering does not recognize.

Related errors

• Unsupported transforms for ACC: {acc_transforms}.
• Unsupported transforms for LHS: {a_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/c53757dafdb30283. Report an issue: GitHub.