This is an automated email from the ASF dual-hosted git repository.
tlopex 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 262a564485 [Fix][Relax][ONNX] Recover ConstantOfShape initializer
shape (#20002)
262a564485 is described below
commit 262a5644858b9d92cd0906f3ee6cdb3dd5f6f03b
Author: Vic Wen <[email protected]>
AuthorDate: Tue Jul 14 20:44:56 2026 +0800
[Fix][Relax][ONNX] Recover ConstantOfShape initializer shape (#20002)
`ConstantOfShape` uses its input tensor as shape metadata. When that
input is an initializer and `keep_params_in_input=True`, the Relax ONNX
frontend should recover the initializer value from `params` instead of
treating the input as an opaque runtime value.
This patch applies `get_constant` to the `ConstantOfShape` shape input
before shape handling. It also guards the constant-shape folding path so
it only calls `len(shape)` on `relax.ShapeExpr` values.
The regression test covers an initializer-backed shape input imported
with `keep_params_in_input=True` and checks that the resulting Relax
function has the expected output shape and dtype.
A separate lint follow-up commit removes stale `F821` suppressions from
two DLight files so the repository-wide CI lint is clean.
Verification:
- `python -m pytest
tests/python/relax/test_frontend_onnx.py::test_constantofshape_initializer_shape_with_keep_params_in_input
-q`
- `pre-commit run --all-files`
Fixes #20001.
---------
Signed-off-by: viiccwen <[email protected]>
---
python/tvm/relax/frontend/onnx/onnx_frontend.py | 8 ++++++--
python/tvm/s_tir/dlight/gpu/gemv.py | 2 +-
python/tvm/s_tir/dlight/gpu/low_batch_gemv.py | 2 +-
tests/python/relax/test_frontend_onnx.py | 23 +++++++++++++++++++++++
4 files changed, 31 insertions(+), 4 deletions(-)
diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py
b/python/tvm/relax/frontend/onnx/onnx_frontend.py
index 0ca7ef9e0c..7f904c8a1b 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ -2124,7 +2124,7 @@ class ConstantOfShape(OnnxOpConverter):
@classmethod
def _impl_v9(cls, bb, inputs, attr, params):
- shape = inputs[0]
+ shape = get_constant(inputs[0], params)
# ONNX spec: `value` is optional and defaults to a zero float32 scalar.
# `get_numpy` requires a TensorProto, so dispatch on presence first.
attr_value = attr.get("value")
@@ -2138,7 +2138,11 @@ class ConstantOfShape(OnnxOpConverter):
shape = relax.ShapeExpr(list(shape.data.numpy()))
# Special case where requested shape are constant
- if len(shape) == 1 and all([isinstance(x, tirx.IntImm) for x in
shape]):
+ if (
+ isinstance(shape, relax.ShapeExpr)
+ and len(shape) == 1
+ and all(isinstance(x, tirx.IntImm) for x in shape)
+ ):
shape = [int(x) for x in shape]
return relax.const(_np.full(shape, value, dtype), dtype)
diff --git a/python/tvm/s_tir/dlight/gpu/gemv.py
b/python/tvm/s_tir/dlight/gpu/gemv.py
index c893c3449a..1c451b964f 100644
--- a/python/tvm/s_tir/dlight/gpu/gemv.py
+++ b/python/tvm/s_tir/dlight/gpu/gemv.py
@@ -14,7 +14,7 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# ruff: noqa: E741, F821
+# ruff: noqa: E741
"""A rule for GEMV and DecodeGEMV."""
from functools import reduce
diff --git a/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py
b/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py
index e854ec1e26..8dedbec4b6 100644
--- a/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py
+++ b/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py
@@ -14,7 +14,7 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# ruff: noqa: E741, F821
+# ruff: noqa: E741
"""A rule for low-batch GEMM / decode-GEMM using GEMV schedule."""
from functools import reduce
diff --git a/tests/python/relax/test_frontend_onnx.py
b/tests/python/relax/test_frontend_onnx.py
index a21b3d4e7a..30d9c35a6c 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -5803,6 +5803,29 @@ def test_constantofshape_default_value():
tvm.ir.assert_structural_equal(tvm_model, Expected)
+def test_constantofshape_initializer_shape_with_keep_params_in_input():
+ shape_init = helper.make_tensor("shape", TensorProto.INT64, [1], [3])
+ node = helper.make_node(
+ "ConstantOfShape",
+ ["shape"],
+ ["y"],
+ value=helper.make_tensor("value", TensorProto.INT64, [1], [1]),
+ )
+ graph = helper.make_graph(
+ [node],
+ "constantofshape_initializer_shape_test",
+ inputs=[helper.make_tensor_value_info("shape", TensorProto.INT64,
[1])],
+ outputs=[helper.make_tensor_value_info("y", TensorProto.INT64, [3])],
+ initializer=[shape_init],
+ )
+ model = helper.make_model(graph,
producer_name="constantofshape_initializer_shape_test")
+
+ tvm_model = from_onnx(model, keep_params_in_input=True)
+
+ assert tuple(dim.value for dim in tvm_model["main"].ret_ty.shape.values)
== (3,)
+ assert tvm_model["main"].ret_ty.dtype == "int64"
+
+
def test_slice():
def verify_slice(data_shape, output_shape, starts, ends, expected,
axes=None, steps=None):
if isinstance(starts, list):
