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

Aperta Adatta ai principianti
  • Qualcosa non va?
#20,948 0 commenti 0 reazioni 0 assegnatari Vedi su GitHub

I maintainer di solito rispondono entro 1 giorno

Nessuno ha ancora preso questa issue.

Valutazione

Difficoltà
1/5
Tempo stimato
Meno di un'ora
Idoneità per principianti
88/100
Tipo di issue
Bug
Chiarezza
Specificata chiaramente
Stato di attività
Attiva
Stack tecnologico
python, pytorch
Ambito
compilers

Direzione di ricerca

La ricerca che non riesce è convert_map in python/tvm/relax/frontend/torch/exported_program_translator.py, che registra prod.default, ma non prod.dim_int. Il metodo _prod esistente in base_fx_graph_translator.py gestisce già gli argomenti dim e keepdim, quindi la modifica consiste nel registrare l’overload ATen mancante in modo che punti a questo metodo. Il lavoro è completato quando lo script di riproduzione nell’issue viene importato senza l’AssertionError e l’output corrisponde a PyTorch.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Descrizione

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
Lingua principale
Python
Stelle
13.9k
Fork
4k
Merge medio
8h 29m
PR unite (30g)
103

Preparare l'ambiente

Questo progetto non fornisce container di sviluppo, Dockerfile né guida per i contributori, quindi l'ambiente è a tuo carico: parti dal suo README e consulta la nostra guida al primo contributo per i passaggi generali.

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di apache/tvm

Tutte le issue di apache/tvm

Issue simili

Altre issue su Python

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.