This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new a5b840d38b [TEST][ONNX] Declare captured symbolic shape variables
explicitly (#20411)
a5b840d38b is described below
commit a5b840d38b3a1c3e0ac5f205b8435254bc922170
Author: Tianqi Chen <[email protected]>
AuthorDate: Tue Sep 22 18:02:23 2026 -0400
[TEST][ONNX] Declare captured symbolic shape variables explicitly (#20411)
Use inline symbolic shapes in the Min/Max expected function signatures
and explicit int64 variables for generated reduction, split, and tile
dimensions. Preserve shared dimensions across input/output annotations
and broadcast expressions.
---
tests/python/relax/test_frontend_onnx.py | 24 ++++++++++++------------
1 file changed, 12 insertions(+), 12 deletions(-)
diff --git a/tests/python/relax/test_frontend_onnx.py
b/tests/python/relax/test_frontend_onnx.py
index d7aa987c05..0efb8030a6 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -40,7 +40,7 @@ from onnx import ModelProto, TensorProto, helper, numpy_helper
import tvm
import tvm.testing
-from tvm import relax
+from tvm import relax, tirx
from tvm.relax.frontend.onnx import from_onnx
from tvm.script import ir as I
from tvm.script import relax as R
@@ -1053,15 +1053,14 @@ def _make_expected_broadcast_ir_min(
Returns:
Expected IR module for the Min operation.
"""
- output_shape = (x_shape[0], 4)
@I.ir_module
class ExpectedMin:
@R.function
def main(
- x: R.Tensor(x_shape, dtype="float32"),
- y: R.Tensor(y_shape, dtype="float32"),
- ) -> R.Tensor(output_shape, dtype="float32"):
+ x: R.Tensor(("n", x_shape[1]), dtype="float32"),
+ y: R.Tensor(("n", y_shape[1]), dtype="float32"),
+ ) -> R.Tensor(("n", 4), dtype="float32"):
n = T.int64()
R.func_attr({"num_input": 2})
with R.dataflow():
@@ -1088,15 +1087,14 @@ def _make_expected_broadcast_ir_max(
Returns:
Expected IR module for the Max operation.
"""
- output_shape = (x_shape[0], 4)
@I.ir_module
class ExpectedMax:
@R.function
def main(
- x: R.Tensor(x_shape, dtype="float32"),
- y: R.Tensor(y_shape, dtype="float32"),
- ) -> R.Tensor(output_shape, dtype="float32"):
+ x: R.Tensor(("n", x_shape[1]), dtype="float32"),
+ y: R.Tensor(("n", y_shape[1]), dtype="float32"),
+ ) -> R.Tensor(("n", 4), dtype="float32"):
n = T.int64()
R.func_attr({"num_input": 2})
with R.dataflow():
@@ -6846,7 +6844,7 @@ def _make_reduce_expected_ir(
def expected_input_shape(shape):
if not dynamic:
return tuple(shape)
- return tuple(f"reduce_dim_{i}" for i in range(len(shape)))
+ return tuple(tirx.Var(f"reduce_dim_{i}", "int64") for i in
range(len(shape)))
axis = None if not axes else tuple(axes)
parser_vars = {
@@ -8784,7 +8782,7 @@ def test_split():
shape = shape_tuple(shape)
if not dynamic:
return shape
- return tuple(f"split_input_dim_{i}" for i in range(len(shape)))
+ return tuple(tirx.Var(f"split_input_dim_{i}", "int64") for i in
range(len(shape)))
dtype = np.dtype(fp_arith).name
input_shape = expected_input_shape(indata_shape)
@@ -9078,7 +9076,9 @@ def test_tile_dynamic_repeats():
def make_expected(dynamic_input, in_shape):
rank = len(in_shape)
input_shape = (
- tuple(f"tile_data_dim_{i}" for i in range(rank)) if dynamic_input
else tuple(in_shape)
+ tuple(tirx.Var(f"tile_data_dim_{i}", "int64") for i in range(rank))
+ if dynamic_input
+ else tuple(in_shape)
)
if rank == 2:
