jax-ml/jax · error · NotImplementedError

Unsupported transforms for LHS: {a_transforms}.

Error message

Unsupported transforms for LHS: {a_transforms}.

What it means

At lowering, the LHS ref's transform stack must match one of the supported patterns: (unswizzle + untiling (+ transpose)) or empty for TMEM refs. Anything else raises NotImplementedError.

Source

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

    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
      case (
          gpu_core.UnswizzleRef(lhs_swizzle),
          gpu_core.UntilingTransform(lhs_tiling),
          state_types.TransposeTransform((1, 0)),
      ):
        lhs_transpose = True
      case () if isinstance(a_ref, tcgen05.TMEMRef):
        lhs_tiling = None
      case _:
        raise NotImplementedError(
            f"Unsupported transforms for LHS: {a_transforms}."
        )
    if not isinstance(a_ref, tcgen05.TMEMRef):
      assert lhs_swizzle is not None
      swizzle_elems = 8 * lhs_swizzle // dtypes.itemsize_bits(a_dtype)
      if lhs_tiling != (8, swizzle_elems):
        raise ValueError("MMA lhs tiling does not fit swizzle. "
                        f"{lhs_tiling=} expected={(8, swizzle_elems)}")

  assert b_transforms_tree is not None
  b_transforms = b_transforms_tree.unflatten(b_transforms_leaves)
  b_out_ty = state_types.transform_type(b_transforms, b_aval)
  assert isinstance(b_out_ty, state_types.AbstractRef)
  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)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use the standard swizzle helper (e.g. swizzle=128) that produces UnswizzleRef+UntilingTransform
2. Transpose via TransposeTransform((1,0)) if a transposed LHS is needed
3. For TMEM LHS, pass the ref with no transforms

Example fix

# before
a_weird = a_ref.transpose((1,0)).disjoint(...)  # unsupported stack
tcgen05.mma(a_weird, b, acc, k_dim=k)
# after
a_smem = load_to_smem(a, swizzle=128)
tcgen05.mma(a_smem, b, acc, k_dim=k)

allowed = lambda t: t in (
    (gpu_core.UnswizzleRef(swz), gpu_core.UntilingTransform(tl), state_types.TransposeTransform((1,0))),
    (gpu_core.UnswizzleRef(swz), gpu_core.UntilingTransform(tl)),
    ())
assert a_transforms in allowed or isinstance(a.ref if hasattr(a,'ref') else a, tcgen05.TMEMRef)

def lhs_transforms_supported(a_transforms):
    for t in a_transforms:
        if not isinstance(t, (gpu_core.UnswizzleRef, gpu_core.UntilingTransform, state_types.TransposeTransform)):
            return False
    return True

When it happens

Trigger: Attaching an arbitrary/disjoint transform or a non-standard swizzle combination to the LHS ref of tcgen05.mma.

Common situations: Custom swizzle helpers composing transforms in a different order; using a DisjointTransform or sliced LHS ref.

Related errors

• Unsupported transforms for ACC: {acc_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/f557efe7e0fcab49. Report an issue: GitHub.