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 ba24a42400 [Relax][Frontend][ONNX] Support Shape start and end
attributes (#20050)
ba24a42400 is described below
commit ba24a424009b841513c7a225a775c2bc3815eccc
Author: Ronald Nap <[email protected]>
AuthorDate: Sun Jul 26 09:47:26 2026 -0700
[Relax][Frontend][ONNX] Support Shape start and end attributes (#20050)
## Summary
This adds support for the `start` and `end` attributes introduced for
ONNX `Shape` in opset 15.
The Relax ONNX frontend previously reused the opset 13 implementation,
which always returned the full input shape. As a result, models using
sliced shape values could construct an incorrect target shape and fail
in downstream operators such as `Reshape`:
```text
ValueError: Reshape expects the new shape to be convertible from the old
shape. However, the old shape is R.shape([12]), with product T.int64(12), while
the new shape is R.shape([2, 3, 4]), with product T.int64(24)
```
### Minimal reproduce
```python
import onnx
from tvm.relax.frontend.onnx import from_onnx
input_shape = [2, 3, 4]
data_shape = [12]
expected_shape = [3, 4]
start = 1
end = None
opset = 15
shape_attrs = {"start": start}
if end is not None:
shape_attrs["end"] = end
model = onnx.helper.make_model(
onnx.helper.make_graph(
[
onnx.helper.make_node("Shape", ["x"], ["shape"], **shape_attrs),
onnx.helper.make_node("Reshape", ["data", "shape"], ["y"]),
],
"shape_start_end_repro",
[
onnx.helper.make_tensor_value_info(
"x", onnx.TensorProto.FLOAT, input_shape
),
onnx.helper.make_tensor_value_info(
"data", onnx.TensorProto.FLOAT, data_shape
),
],
[
onnx.helper.make_tensor_value_info(
"y", onnx.TensorProto.FLOAT, expected_shape
)
],
),
opset_imports=[onnx.helper.make_opsetid("", opset)],
)
print(f"Shape attributes: start={start}, end={end}")
print(f"Expected Shape output: {input_shape[start:end]}")
print(from_onnx(model, opset=opset).script())
```
The new implementation applies `start` and `end` slicing to static and
symbolic shape expressions. It also handles runtime-produced shape
values by converting them to a tensor, applying `strided_slice`, and
converting the result back to a shape.
---
python/tvm/relax/frontend/onnx/onnx_frontend.py | 27 +++-
tests/python/relax/test_frontend_onnx.py | 170 ++++++++++++++++++++++++
2 files changed, 196 insertions(+), 1 deletion(-)
diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py
b/python/tvm/relax/frontend/onnx/onnx_frontend.py
index 881789c7ed..806d16f5a8 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ -1629,7 +1629,7 @@ class Clip(OnnxOpConverter):
class Shape(OnnxOpConverter):
- """Converts an onnx Equal node into an equivalent Relax expression."""
+ """Converts an onnx Shape node into an equivalent Relax expression."""
@classmethod
def _impl_v13(cls, bb, inputs, attr, params):
@@ -1647,6 +1647,31 @@ class Shape(OnnxOpConverter):
return data_info.shape
+ @classmethod
+ def _impl_v15(cls, bb, inputs, attr, params):
+ shape = cls._impl_v13(bb, inputs, attr, params)
+ start = attr.get("start", 0)
+ end = attr.get("end")
+
+ if start == 0 and end is None:
+ return shape
+
+ if isinstance(shape, relax.ShapeExpr):
+ return relax.ShapeExpr(list(shape.values)[start:end])
+
+ shape_tensor = bb.normalize(relax.op.shape_to_tensor(shape))
+ sliced_shape = bb.normalize(
+ relax.op.strided_slice(
+ shape_tensor,
+ axes=[0],
+ begin=[start],
+ end=[end if end is not None else 2**63 - 1],
+ strides=[1],
+ assume_inbound=False,
+ )
+ )
+ return bb.normalize(relax.op.tensor_to_shape(sliced_shape))
+
class Trilu(OnnxOpConverter):
"""Given a 2-D matrix or batches of 2-D matrices, returns the upper or
diff --git a/tests/python/relax/test_frontend_onnx.py
b/tests/python/relax/test_frontend_onnx.py
index 7e3da83df0..3a0a4aa5b9 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -3067,6 +3067,176 @@ def test_shape():
tvm.ir.assert_structural_equal(tvm_model, Expected)
[email protected](
+ "attrs,expected_shape",
+ [
+ ({"start": 1}, (4, 5, 6)),
+ ({"end": -1}, (3, 4, 5)),
+ ({"start": -2}, (5, 6)),
+ ({"start": 1, "end": 3}, (4, 5)),
+ ({"start": -10, "end": 10}, (3, 4, 5, 6)),
+ ({"start": 3, "end": 2}, ()),
+ ],
+)
+def test_shape_start_end(attrs, expected_shape):
+ expected_shape = list(expected_shape)
+ shape_node = helper.make_node("Shape", ["data"], ["output"], **attrs)
+
+ graph = helper.make_graph(
+ [shape_node],
+ "shape_start_end_test",
+ inputs=[
+ helper.make_tensor_value_info(
+ "data",
+ TensorProto.FLOAT,
+ [3, 4, 5, 6],
+ ),
+ ],
+ outputs=[
+ helper.make_tensor_value_info(
+ "output",
+ TensorProto.INT64,
+ [len(expected_shape)],
+ )
+ ],
+ )
+
+ model = helper.make_model(
+ graph,
+ producer_name="shape_start_end_test",
+ opset_imports=[helper.make_opsetid("", 15)],
+ )
+ tvm_model = from_onnx(
+ model,
+ opset=15,
+ keep_params_in_input=True,
+ )
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(
+ data: R.Tensor((3, 4, 5, 6), dtype="float32"),
+ ) -> R.Shape(expected_shape):
+ R.func_attr({"num_input": 1})
+ with R.dataflow():
+ gv: R.Shape(expected_shape) = R.shape(expected_shape)
+ R.output(gv)
+ return gv
+
+ tvm.ir.assert_structural_equal(tvm_model, Expected)
+
+
+def test_shape_start_end_symbolic():
+ shape_node = helper.make_node(
+ "Shape",
+ ["data"],
+ ["output"],
+ start=1,
+ end=3,
+ )
+ graph = helper.make_graph(
+ [shape_node],
+ "shape_start_end_symbolic_test",
+ inputs=[
+ helper.make_tensor_value_info(
+ "data",
+ TensorProto.FLOAT,
+ [3, "B", 5, 6],
+ ),
+ ],
+ outputs=[
+ helper.make_tensor_value_info(
+ "output",
+ TensorProto.INT64,
+ [2],
+ )
+ ],
+ )
+
+ model = helper.make_model(
+ graph,
+ producer_name="shape_start_end_symbolic_test",
+ opset_imports=[helper.make_opsetid("", 15)],
+ )
+ tvm_model = from_onnx(
+ model,
+ opset=15,
+ keep_params_in_input=True,
+ )
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(
+ data: R.Tensor((3, "B", 5, 6), dtype="float32"),
+ ) -> R.Shape(ndim=2):
+ B = T.int64()
+ R.func_attr({"num_input": 1})
+ with R.dataflow():
+ gv: R.Shape([B, 5]) = R.shape([B, 5])
+ R.output(gv)
+ return gv
+
+ tvm.ir.assert_structural_equal(tvm_model, Expected)
+
+
+def test_shape_start_end_scalar():
+ shape_node = helper.make_node(
+ "Shape",
+ ["data"],
+ ["output"],
+ start=1,
+ )
+
+ graph = helper.make_graph(
+ [shape_node],
+ "shape_start_end_scalar_test",
+ inputs=[
+ helper.make_tensor_value_info(
+ "data",
+ TensorProto.FLOAT,
+ [],
+ ),
+ ],
+ outputs=[
+ helper.make_tensor_value_info(
+ "output",
+ TensorProto.INT64,
+ [0],
+ )
+ ],
+ )
+
+ model = helper.make_model(
+ graph,
+ producer_name="shape_start_end_scalar_test",
+ opset_imports=[helper.make_opsetid("", 15)],
+ )
+ tvm_model = from_onnx(
+ model,
+ opset=15,
+ keep_params_in_input=True,
+ )
+
+ assert relax.analysis.check_well_formed(tvm_model)
+
+ op_names = []
+
+ def collect_ops(expr):
+ if isinstance(expr, relax.Call) and isinstance(expr.op, tvm.ir.Op):
+ op_names.append(expr.op.name)
+
+ relax.analysis.post_order_visit(tvm_model["main"], collect_ops)
+
+ assert op_names == [
+ "relax.shape_of",
+ "relax.shape_to_tensor",
+ "relax.strided_slice",
+ "relax.tensor_to_shape",
+ ]
+
+
def test_trilu():
def verify_trilu(upper: bool):
node = helper.make_node("Trilu", ["x"], ["y"], upper=upper)
