[Bug][Relax] ReorderTakeAfterMatmul fails on vector dot products

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

Los mantenedores suelen responder en 1 día

@hiyufan ya está trabajando en esto.

Desde el 25/10/2026.

  • #20815 de @hiyufan — abierto

Evaluación

Dificultad
2/5
Tiempo estimado
1-3 horas
Aptitud para principiantes
74/100
Tipo de issue
Error
Claridad
Bien especificado
Estado de actividad
Activo
Stack tecnológico
cpp, python
Área
compilers

Línea de trabajo

Empieza en src/relax/transform/reorder_take_after_matmul.cc, donde el issue identifica como puntos de fallo la condición de reescritura para rango uno y el eje take de reemplazo. Reproduce el problema con el módulo Relax de producto escalar de vectores proporcionado y, después, comprueba el caso de control adyacente de rango dos. Se considera terminado cuando el pass conserva el resultado escalar o deja sin cambios el patrón de rango no compatible, sin generar un error.

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

Descripción

Description

ReorderTakeAfterMatmul rewrites a valid vector dot product into a scalar MatMul followed by take(axis=-1). A scalar has no valid axis, so the pass raises while constructing the replacement.

Environment

  • TVM 0.26.dev0, executed at e269315c90e3a061c9e1c77b370ce883b1b223f4
  • Ubuntu 24.04 x86-64, LLVM 18.1.3 CPU target
  • Current upstream main checked at d55759ec1ebed554d9782a52e43e5730b8a96656
  • src/relax/transform/reorder_take_after_matmul.cc is byte-identical at both revisions (SHA-256 80324bd216442834d9c181f2466fc6ff8aa5538b93e78251b59f863ce5a6e6ba)
  • I did not have a binary built from current main, so the runtime reproduction below is from the executed revision

Minimal reproduction

from tvm import relax
from tvm.script import ir as I
from tvm.script import relax as R

@I.ir_module
class Module:
    @R.function
    def main(
        lhs: R.Tensor((3,), "float32"),
        weight: R.Tensor((3,), "float32"),
        indices: R.Tensor((3,), "int64"),
    ):
        selected = R.take(weight, indices, axis=0, mode="clip")
        return R.matmul(lhs, selected)

assert relax.analysis.check_well_formed(Module, check_ty=True)
relax.transform.ReorderTakeAfterMatmul()(Module)

Expected behavior

The pass should preserve the result or leave this unsupported rank pattern unchanged.

Actual behavior

The source is well formed, executes on LLVM, and returns a scalar. The pass raises:

ValueError: the input axis -1 is out of range. The input tensor has 0 dimensions

The adjacent rank-two weight control is rewritten successfully and its source and target outputs match.

Likely cause

The rewrite guard accepts a rank-one weight and indices tensor. The replacement MatMul therefore produces a scalar, but the pass constructs take(out_table, indices, matmul_ty->ndim - 1, mode), yielding axis -1 for a rank-zero value.

Duplicate check

I searched open and closed TVM issues and PRs for ReorderTakeAfterMatmul, vector weights, scalar outputs, and axis -1. #20198/#20331 concern out_dtype, and #20201/#20206 concern take mode; I did not find this rank-guard failure.

I can send a fix PR and help with follow-up testing if this diagnosis looks right.

Triage

  • needs-triage
  • type: bug
Lenguaje dominante
Python
Estrellas
13.9k
Forks
4k
Merge medio
8 h 2 min
PR fusionados (30 d)
180

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.