jax-ml/jax · error · ValueError

TMEM transpose not allowed.

Error message

TMEM transpose not allowed.

What it means

The tcgen05 MMA lowering cannot transpose an operand that lives in Tensor Core Memory (TMEM). When lhs_transpose=True and the lhs is a TMEMRef, Mosaic GPU raises this error because TMEM has no transposing memref transform.

Source

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

        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:
    lhs_swizzle = rhs_swizzle
  elif rhs_swizzle != lhs_swizzle:
    raise ValueError("MMA rhs swizzle must match lhs swizzle."
                      f" {lhs_swizzle=} {rhs_swizzle=}")
  if lhs_transpose:
    if isinstance(a_ref, tcgen05.TMEMRef):
      raise ValueError("TMEM transpose not allowed.")
    a_ref = mgpu.memref_transpose(a_ref, (1, 0, 3, 2))
  if rhs_transpose:
    b_ref = mgpu.memref_transpose(b_ref, (1, 0, 3, 2))
  if isinstance(accumulate, bool):
    accumulate = mgpu.c(accumulate, ir.IntegerType.get_signless(1))
  elif isinstance(accumulate, mgpu.FragmentedArray):
    accumulate = accumulate.registers.item()
    assert isinstance(accumulate, ir.Value)

  if a_scale_ref is not None and a_scale_transforms_tree is not None:
    assert isinstance(a_scale_ref_aval, state.AbstractRef)
    a_scale_transforms = a_scale_transforms_tree.unflatten(
        a_scale_transforms_leaves
    )
    a_scale_transform_avals = a_scale_transforms_tree.unflatten(
        a_scale_transforms_leaves_avals
    )
    a_scale_ref, _, a_scale_transforms = lowering._handle_transforms(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Move the operand to SMEM (a regular memref) before transposing, then transpose
2. Pre-transpose the data before loading into TMEM (transpose during the SMEM-to-TMEM copy)

Example fix

// before
acc = tcgen05_dot(...)
tcgen05_mma(lhs_tmem_ref, rhs, acc, lhs_transpose=True)
// after
lhs_smem = ...  # copy from TMEM to SMEM first
tcgen05_mma(mgpu.memref_transpose(lhs_smem, (1,0,3,2)), rhs, acc)

assert not (lhs_transpose and isinstance(lhs, tcgen05.TMEMRef))

def can_transpose(ref): return not isinstance(ref, tcgen05.TMEMRef)

When it happens

Trigger: Passing a tcgen05.TMEMRef as the lhs (A) operand together with lhs_transpose=True to tcgen05_mma.

Common situations: Reusing a TMEM accumulator/operand produced by a previous MMA or async_copy_to_tmem and then trying to transpose it in a follow-up MMA; porting SMEM-based kernels to TMEM pipelines.

Related errors

• Can't transpose a TMEM reference.
• TMEM aliasing only supported for Refs with the same first di
• Unsupported TMEM ref {ref}.
• Stores to TMEM are asynchronous operations and cannot be per
• Accumulator must be a TMEM Ref.

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