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

未关闭 适合新手
  • 哪里不对？
#20,948 0 条评论 0 个 reaction 已指派 0 人 在 GitHub 查看

维护者通常 1 天内回复

还没有人认领这个 Issue。

评估

难度
1/5
预计耗时
1 小时以内
新手友好度
88/100
Issue 类型
缺陷
描述清晰度
描述清楚
活跃度
活跃
技术栈
python, pytorch
领域
compilers

调研方向

失败的查找项是 python/tvm/relax/frontend/torch/exported_program_translator.py 中的 convert_map，它注册了 prod.default，但没有注册 prod.dim_int。base_fx_graph_translator.py 中现有的 _prod 方法已经处理了 dim 和 keepdim 参数，因此需要做的更改是注册缺失的 ATen 重载，使其指向该方法。完成标准是 issue 中的复现脚本导入时不再出现 AssertionError，并且输出与 PyTorch 一致。

由索引模型根据 Issue 内容生成。

描述

needs-triage type: bug
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
主要语言
Python
星标
13.9k
派生
4k
平均合并
8 小时 29 分钟
30 天内合并 PR
103

环境准备

这个项目没有提供开发容器、Dockerfile 或贡献指南，环境需要你自己搭建：先看它的 README，通用步骤见我们的新手贡献指南。

从这里开始

  1. 先读完整个 Issue，再读项目的贡献指南。
  2. 在 Issue 下留言说明你要接手 —— 这能避免两个人做同样的事。
  3. Fork 仓库，在一个分支上完成修改。
  4. 提交 Pull Request，并在描述里引用这个 Issue 编号。

apache/tvm 的其他 Issue

查看 apache/tvm 的全部 Issue

相似的 Issue

更多 Python Issue

把新 issue 发到你的邮箱

精选适合新手参与的 GitHub issue 摘要。