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

オープン 初心者向け
  • 何か問題がありますか？
#20,948 コメント 0 件 リアクション 0 件 担当者 0 名 GitHub で見る

メンテナーはふだん 1 日以内に返信

まだ誰も着手していません。

評価

難易度
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分
マージ済み PR（30日）
103

環境構築

このプロジェクトには開発コンテナ、Dockerfile、コントリビューションガイドがありません。まず README を読み、一般的な手順ははじめてのコントリビューションガイドを参照してください。

はじめの一歩

  1. issue を最後まで読み、次にプロジェクトのコントリビューションガイドを読みます。
  2. 着手することを issue にコメントします — 二人が同じ作業をするのを防げます。
  3. リポジトリをフォークし、ブランチを切って変更します。
  4. issue 番号を参照したプルリクエストを送ります。

apache/tvm のほかの issue

apache/tvm の issue をすべて見る

似ている issue

Python の issue をもっと見る

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。