{"record":{"id":"da80368b4ee96237","repo":"jax-ml/jax","slug":"tmem-transpose-not-allowed-da8036","errorCode":null,"errorMessage":"TMEM transpose not allowed.","messagePattern":"TMEM transpose not allowed\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":2781,"sourceCode":"        f\" {rhs_tiling=} expected={(8, swizzle_elems)}\"\n    )\n\n  if barrier_transforms_tree is not None and barrier_ref is not None:\n    barrier_transforms = barrier_transforms_tree.unflatten(\n        barrier_transforms_leaves\n    )\n    base_index = _get_barrier_base_index(barrier_ref_aval, barrier_transforms)\n    if base_index is not None:\n      barrier_ref = barrier_ref[base_index]\n\n  if lhs_swizzle is None:\n    lhs_swizzle = rhs_swizzle\n  elif rhs_swizzle != lhs_swizzle:\n    raise ValueError(\"MMA rhs swizzle must match lhs swizzle.\"\n                      f\" {lhs_swizzle=} {rhs_swizzle=}\")\n  if lhs_transpose:\n    if isinstance(a_ref, tcgen05.TMEMRef):\n      raise ValueError(\"TMEM transpose not allowed.\")\n    a_ref = mgpu.memref_transpose(a_ref, (1, 0, 3, 2))\n  if rhs_transpose:\n    b_ref = mgpu.memref_transpose(b_ref, (1, 0, 3, 2))\n  if isinstance(accumulate, bool):\n    accumulate = mgpu.c(accumulate, ir.IntegerType.get_signless(1))\n  elif isinstance(accumulate, mgpu.FragmentedArray):\n    accumulate = accumulate.registers.item()\n    assert isinstance(accumulate, ir.Value)\n\n  if a_scale_ref is not None and a_scale_transforms_tree is not None:\n    assert isinstance(a_scale_ref_aval, state.AbstractRef)\n    a_scale_transforms = a_scale_transforms_tree.unflatten(\n        a_scale_transforms_leaves\n    )\n    a_scale_transform_avals = a_scale_transforms_tree.unflatten(\n        a_scale_transforms_leaves_avals\n    )\n    a_scale_ref, _, a_scale_transforms = lowering._handle_transforms(","sourceCodeStart":2763,"sourceCodeEnd":2799,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L2763-L2799","documentation":"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.","triggerScenarios":"Passing a tcgen05.TMEMRef as the lhs (A) operand together with lhs_transpose=True to tcgen05_mma.","commonSituations":"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.","solutions":["Move the operand to SMEM (a regular memref) before transposing, then transpose","Pre-transpose the data before loading into TMEM (transpose during the SMEM-to-TMEM copy)"],"exampleFix":"// before\nacc = tcgen05_dot(...)\ntcgen05_mma(lhs_tmem_ref, rhs, acc, lhs_transpose=True)\n// after\nlhs_smem = ...  # copy from TMEM to SMEM first\ntcgen05_mma(mgpu.memref_transpose(lhs_smem, (1,0,3,2)), rhs, acc)","handlingStrategy":"type-guard","validationCode":"assert not (lhs_transpose and isinstance(lhs, tcgen05.TMEMRef))","typeGuard":"def can_transpose(ref): return not isinstance(ref, tcgen05.TMEMRef)","tryCatchPattern":null,"preventionTips":["Track which refs are TMEM vs SMEM in kernel structure","Transposes belong in SMEM, before TMEM stores"],"tags":["jax","pallas","tcgen05","tmem","transpose"],"backgroundTag":"unsupported-operation-on-memory-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}