[Feature Request][Relax][Torch] Add support for torch.mm (mm.default) in from_exported_program

[tinywisdom]

Expected behavior

• The frontend should support lowering torch.mm into an equivalent Relax operator (likely relax.op.matmul).

• At minimum, if unsupported, provide a clearer user-facing error message and possible fallback guidance.

Actual behavior

When converting a PyTorch torch.exported program into TVM Relax using from_exported_program, a model containing torch.mm fails with:

AssertionError: Unsupported function types ['mm.default']

This indicates that the mm.default operator (2D matrix multiply) is currently not supported in the TVM Relax PyTorch frontend.

Environment

• OS: (Ubuntu 22.04.4 LTS (x86_64))
• TVM version: (release v0.21.0)
• Python: (3.10.16)
• LLVM: (17.0.6)

Steps to reproduce

import torch
import torch.nn as nn
from torch.export import export as torch_export
from tvm.relax.frontend.torch import from_exported_program

class M(nn.Module):
    def forward(self, a, b):
        # Key op: torch.mm (2D matrix multiply)
        return torch.mm(a, b)

def main():
    torch.manual_seed(0)
    m = M().eval()

    # Inputs: (2, 3) @ (3, 4) -> (2, 4)
    a = torch.randn(2, 3, dtype=torch.float32)
    b = torch.randn(3, 4, dtype=torch.float32)

    # 1) Check eager path
    with torch.inference_mode():
        y = m(a, b)
    print("PyTorch eager OK, y.shape =", tuple(y.shape))

    # 2) Export
    ep = torch_export(m, (a, b))
    print("ExportedProgram created.")

    # 3) Import into TVM Relax — triggers unsupported function type
    mod = from_exported_program(ep)

if __name__ == "__main__":
    main()

Triage

• needs-triage
• bug

cc @junrushao @shingjan

[MasterJH5574]

Completed with #18343 merged in. Thanks @tlopex for working on this!