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

Abierto Apto para principiantes
  • ¿Algo no está bien?
#20,948 0 comentarios 0 reacciones 0 asignados Ver en GitHub

Los mantenedores suelen responder en 1 día

Nadie ha tomado este issue todavía.

Evaluación

Dificultad
1/5
Tiempo estimado
Menos de una hora
Aptitud para principiantes
88/100
Tipo de issue
Error
Claridad
Bien especificado
Estado de actividad
Activo
Stack tecnológico
python, pytorch
Área
compilers

Línea de trabajo

La búsqueda que falla es convert_map en python/tvm/relax/frontend/torch/exported_program_translator.py, que registra prod.default, pero no prod.dim_int. El método _prod existente en base_fx_graph_translator.py ya gestiona los argumentos dim y keepdim, así que el cambio consiste en registrar la sobrecarga de ATen que falta para que apunte a ese método. Se considera terminado cuando el script de reproducción del issue se importa sin el AssertionError y la salida coincide con PyTorch.

Escrito por el modelo de indexación a partir del texto del issue.

Descripción

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
Lenguaje dominante
Python
Estrellas
13.9k
Forks
4k
Merge medio
8 h 29 min
PR fusionados (30 d)
103

Preparar el entorno

Este proyecto no incluye contenedor de desarrollo, Dockerfile ni guía de contribución, así que la configuración corre por tu cuenta: empieza por su README y consulta nuestra guía para la primera contribución para los pasos generales.

Primeros pasos

  1. Lee el issue completo y luego la guía de contribución del proyecto.
  2. Comenta en el issue que vas a ocuparte — evita que dos personas hagan lo mismo.
  3. Haz un fork del repositorio y trabaja en una rama.
  4. Abre un pull request que haga referencia al número del issue.

Más de apache/tvm

Todos los issues de apache/tvm

Issues similares

Más issues de Python

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.