[Bug][Relax][PyTorch] torch.prod with reduction dimension fails in from_exported_program: "AssertionError: Unsupported function types ['prod.dim_int']" #20948

Expected behavior

Calling torch.prod(x, dim=...) in a PyTorch module should be successfully imported into Relax via from_exported_program. The underlying translator already has an implementation in base_fx_graph_translator.py (_prod) that supports dim and keepdim.

Actual behavior

When translating an ExportedProgram that uses torch.prod with a reduction axis, from_exported_program fails with an AssertionError:

Traceback (most recent call last):
  File "repro.py", line 11, in <module>
    mod = from_exported_program(ep)
  File "python/tvm/relax/frontend/torch/exported_program_translator.py", line 2454, in from_exported_program
    return ExportedProgramImporter().from_exported_program(...)
  File "python/tvm/relax/frontend/torch/exported_program_translator.py", line 2289, in from_exported_program
    output_args = self._translate_fx_graph(...)
  File "python/tvm/relax/frontend/torch/exported_program_translator.py", line 1530, in _translate_fx_graph
    self._check_unsupported_func_type(unsupported_ops)
  File "python/tvm/relax/frontend/torch/base_fx_graph_translator.py", line 324, in _check_unsupported_func_type
    assert len(unsupported_ops) == 0, f"Unsupported function types {unsupported_ops}"
AssertionError: Unsupported function types ['prod.dim_int']

Environment

• TVM Version: 0.28.dev0 (commit cbb1e12fe on main)
• PyTorch Version: 2.5.1+cpu
• Python Version: 3.12.2
• Operating System: Linux x86_64

Steps to reproduce

import torch
from torch.export import export
from tvm.relax.frontend.torch import from_exported_program

class ProdModel(torch.nn.Module):
    def forward(self, x):
        return torch.prod(x, dim=1)

x = torch.randn(3, 4)
ep = export(ProdModel(), (x,))
mod = from_exported_program(ep)

Root Cause Analysis

In PyTorch ATen:

• Calling torch.prod(x) (global reduction across all axes) emits torch.ops.aten.prod.default(x).
• Calling torch.prod(x, dim=1) emits torch.ops.aten.prod.dim_int(x, 1).

In python/tvm/relax/frontend/torch/base_fx_graph_translator.py, _prod already implements dimension extraction:

def _prod(self, node: fx.Node) -> relax.Var:
    args = self.retrieve_args(node)
    x = args[0]
    dim = args[1] if len(node.args) > 1 else node.kwargs.get("dim", None)
    keepdim = args[2] if len(node.args) > 2 else node.kwargs.get("keepdim", False)
    return self.block_builder.emit(relax.op.prod(x, dim, keepdims=keepdim))

However, in python/tvm/relax/frontend/torch/exported_program_translator.py, only "prod.default" was registered in convert_map:

"prod.default": self._prod,

"prod.dim_int" was omitted from convert_map.

Proposed Fix

In python/tvm/relax/frontend/torch/exported_program_translator.py, add "prod.dim_int": self._prod to convert_map:

diff --git a/python/tvm/relax/frontend/torch/exported_program_translator.py b/python/tvm/relax/frontend/torch/exported_program_translator.py
--- a/python/tvm/relax/frontend/torch/exported_program_translator.py
+++ b/python/tvm/relax/frontend/torch/exported_program_translator.py
@@ -1967,6 +1967,7 @@ class ExportedProgramImporter(BaseFXGraphImporter):
             "pow.Tensor_Tensor": self._pow,
             "pow.Scalar": self._pow,
             "prod.default": self._prod,
+            "prod.dim_int": self._prod,
             "reciprocal.default": self._reciprocal,
             "remainder.Scalar": self._remainder,
             "remainder.Tensor": self._remainder,

Triage

• needs-triage
• relax:frontend