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

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

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

@hiyufan がすでに取り組んでいます。

2026年10月25日 から。

  • #20815 @hiyufan による — オープン

評価

難易度
2/5
見積もり時間
1〜3時間
初心者へのやさしさ
74/100
issue の種類
バグ
明瞭さ
明確に書かれている
活発さ
活発
技術スタック
cpp, python
領域
compilers

調査の方向性

src/relax/transform/reorder_take_after_matmul.ccから始めます。issueでは、ランク1の書き換えガードと置換後のtake軸が失敗箇所として特定されています。提供されたベクトル内積のRelaxモジュールで再現し、その後、隣接するランク2の対照ケースを確認します。passがスカラー結果を保持するか、サポートされていないランクパターンをエラーを発生させずに変更しないままにすれば完了です。

索引モデルが issue の本文から書いたものです。

説明

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
主要言語
Python
スター
13.9k
フォーク
4k
平均マージ
8時間 6分
マージ済み PR（30日）
184

環境構築

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

はじめの一歩

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

apache/tvm のほかの issue

apache/tvm の issue をすべて見る

似ている issue

Python の issue をもっと見る

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

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