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

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()