[Bug] [Relax][ONNX] Expand uses a symbolic requested shape without merging broadcast dimensions

未关闭
#20,897 0 条评论 0 个 reaction 已指派 0 人 在 GitHub 查看

维护者通常 1 天内回复

还没有人认领这个 Issue。

评估

难度
4/5
预计耗时
3-5 天
新手友好度
68/100
Issue 类型
缺陷
描述清晰度
描述清楚
活跃度
活跃
技术栈
python
领域
compilers

调研方向

Start with the ShapeExpr branch in python/tvm/relax/frontend/onnx/onnx_frontend.py, around lines 3453–3515, and trace how it handles the requested shape in Expand. Run the supplied reproduction with z_shape ["N", 2]; it should produce the same [3, 2] output as the constant-shape case, without dropping rows.

由索引模型根据 Issue 内容生成。

描述

needs-triage type: bug
Expected behavior

Expanding [3,1] with requested shape [1,2] should produce [3,2] under ONNX multidirectional broadcasting, including when the requested shape comes from Shape(z).

Actual behavior

With the same actual inputs, declaring z as [1,2] gives the correct [3,2] output. Changing only its annotation to [N,2] silently produces [1,2] and drops two rows:

[1, 2] TVM (3, 2) [[1.0, 1.0], [2.0, 2.0], [3.0, 3.0]]
['N', 2] TVM (1, 2) [[1.0, 1.0]]

The ShapeExpr branch merges only constant target extents. This differs from #20873: the target here is a symbolic ShapeExpr and conversion succeeds.

Environment

TVM 0.28.dev0, LLVM 18.1.3, Linux x86_64; Python 3.12.3, NumPy 2.5.4, ONNX 1.22.0, ORT 1.29.0. The installed binary build revision is unknown; current main was inspected statically, not rebuilt.

Steps to reproduce

Run with an LLVM-enabled TVM installation:

import numpy as np
import onnx
import tvm
from onnx import TensorProto as T, helper as h
from onnx.reference import ReferenceEvaluator
from tvm import relax
from tvm.relax.frontend.onnx import from_onnx

x = np.array([[1], [2], [3]], dtype="float32")
z = np.zeros((1, 2), dtype="float32")
expected = np.repeat(x, 2, axis=1)
for z_shape in ([1, 2], ["N", 2]):
    model = h.make_model(h.make_graph([
        h.make_node("Shape", ["z"], ["shape"]),
        h.make_node("Expand", ["x", "shape"], ["y"])], "expand",
        [h.make_tensor_value_info("x", T.FLOAT, [3, 1]),
         h.make_tensor_value_info("z", T.FLOAT, z_shape)],
        [h.make_tensor_value_info("y", T.FLOAT, [3, 2])]),
        opset_imports=[h.make_opsetid("", 18)], ir_version=10)
    onnx.checker.check_model(model, full_check=True)
    np.testing.assert_array_equal(ReferenceEvaluator(model).run(None, {"x": x, "z": z})[0], expected)
    try:
        mod = from_onnx(model, keep_params_in_input=False)
        vm = relax.VirtualMachine(relax.build(mod, target="llvm"), tvm.cpu())
        actual = vm["main"](tvm.runtime.tensor(x), tvm.runtime.tensor(z)).numpy()
        print(z_shape, "TVM", actual.shape, actual.tolist())
        if z_shape[0] == 1:
            np.testing.assert_array_equal(actual, expected)
    except Exception as error:
        if z_shape[0] == 1:
            raise
        print(z_shape, type(error).__name__, str(error).splitlines()[-1])
Triage
  • needs-triage
主要语言
Python
星标
13.9k
派生
4k
平均合并
8 小时 6 分钟
30 天内合并 PR
147

环境准备

这个项目没有提供开发容器、Dockerfile 或贡献指南,环境需要你自己搭建:先看它的 README,通用步骤见我们的新手贡献指南。

从这里开始

  1. 先读完整个 Issue,再读项目的贡献指南。
  2. 在 Issue 下留言说明你要接手 —— 这能避免两个人做同样的事。
  3. Fork 仓库,在一个分支上完成修改。
  4. 提交 Pull Request,并在描述里引用这个 Issue 编号。

apache/tvm 的其他 Issue

查看 apache/tvm 的全部 Issue

相似的 Issue

更多 Python Issue

把新 issue 发到你的邮箱

精选适合新手参与的 GitHub issue 摘要。