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 b205abd673 [FIX][Relax][ONNX] Keep the static shape of a rank-0 Shape
input (#20092)
b205abd673 is described below
commit b205abd673f02af45d1448c8fd4e463c69512a9d
Author: Aditya Singh <[email protected]>
AuthorDate: Wed Aug 5 16:53:33 2026 -0700
[FIX][Relax][ONNX] Keep the static shape of a rank-0 Shape input (#20092)
`Shape._impl_v13` in the Relax ONNX frontend chose between the static
shape and a runtime `shape_of` with a truthiness test:
```python
if not data_info.shape:
return bb.normalize(relax.op.shape_of(inputs[0]))
return data_info.shape
```
A rank-0 tensor has a defined but empty shape, and an empty `ShapeExpr`
is falsy, so a scalar input took the runtime path that is meant only for
a tensor whose shape is genuinely unknown. `TensorType.shape` is `None`
in that case and `R.shape([])` for a rank-0 tensor, so the two are
distinguishable, and only the first should reach `shape_of`.
The value handed back was then a normalized `R.shape_of` call rather
than a `ShapeExpr`, so every downstream converter that matches on
`relax.ShapeExpr` lost its static path. `Slice` is the visible case from
the report: importing `Shape` followed by `Slice` on a scalar input
raises
```
Error converting operator Slice, with inputs: [R.shape_of(X), ...]
ValueError: Slice requires a statically known input rank.
```
because `_get_known_tensor_rank` cannot produce a rank for a shape
value. `Gather` and `Reshape` carry the same `isinstance(...,
relax.ShapeExpr)` match, so they are exposed to the same shape.
## Fix
Compare against `None`, so a rank-0 input keeps the static `R.shape([])`
that every other rank already gets.
`Shape` of a scalar now folds at import time to `R.shape([])`, and
`Shape` followed by `Slice` folds to an empty `int64` tensor of shape
`(0,)`. Checked against ONNX Runtime 1.24 on the same graph: `Shape` of
a scalar returns an empty `int64` array of shape `(0,)`, and slicing it
returns the same, so the folded result matches.
## Effect on an existing test
`test_shape_start_end_scalar`, added by #20050, pinned the runtime
fallback for a rank-0 input with `start=1`, asserting the op chain
`relax.shape_of, relax.shape_to_tensor, relax.strided_slice,
relax.tensor_to_shape`. With the static shape preserved, that case folds
to the same empty static shape, so the test now asserts the folded
module and that no ops remain. ONNX Runtime returns an empty `int64`
array for that graph too, so the folded answer is the correct one, and
the assertion change is the point of the fix rather than a workaround
for it.
## Testing
Two new tests in `tests/python/relax/test_frontend_onnx.py`:
- `test_shape_scalar_input`, structural equality against the expected
module, pins that `Shape` of a rank-0 input emits `R.shape([])` and not
`R.shape_of`.
- `test_slice_of_scalar_shape`, the reported pattern end to end, pins
that the import succeeds and yields an empty `int64` tensor.
Verified fail-before and pass-after against the base ref rather than a
stash, over the whole file so any collateral damage would show:
```
python -m pytest tests/python/relax/test_frontend_onnx.py -q # with the
fix
python -m pytest tests/python/relax/test_frontend_onnx.py -q # at
upstream/main
```
The failure sets differ by exactly three entries, all of them the target
tests, and there are no new failures:
```
only in the fixed run: (none)
only in the base run: test_shape_scalar_input
test_shape_start_end_scalar
test_slice_of_scalar_shape
```
This was run on a local build configured with `USE_LLVM OFF`, so the
tests that go through `check_correctness` and
`tvm.compile(target="llvm")` fail identically in both runs with
`ValueError: Cannot find global function target.build.llvm`. They are
the same 174 entries on both sides and are unrelated to this change. The
three tests above and the surrounding `test_shape` and
`test_shape_start_end` cases need no codegen and were run directly, 16
passed.
Lint checked with the pinned `ruff==0.12.3` from
`.pre-commit-config.yaml`: `ruff format --check` reports already
formatted and `ruff check` passes on both files.
Fixes #17770
---
python/tvm/relax/frontend/onnx/onnx_frontend.py | 26 ++++-
tests/python/relax/test_frontend_onnx.py | 120 ++++++++++++++++++++++--
2 files changed, 137 insertions(+), 9 deletions(-)
diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py
b/python/tvm/relax/frontend/onnx/onnx_frontend.py
index d9d97126af..c7e6b5b585 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ -222,10 +222,21 @@ def get_info(
Returns
-------
- Tuple[str, List, str, List, Dict]
+ Tuple[str, Optional[List], str, Optional[List], Dict]
The name, shape, type, and shape name of the ValueInfoProto, and the
- value_dict.
+ value_dict. The shape and shape name are None when the proto carries no
+ shape field at all, which means the rank is unknown.
"""
+ tensor_type = info_proto.type.tensor_type
+ elem_dtype = get_type(tensor_type.elem_type) if tensor_type.elem_type else
None
+
+ # An absent shape field means unknown rank, which is not the same as a
rank-0
+ # tensor whose shape field is present with zero dims. Both would otherwise
+ # collapse to an empty list and become R.Tensor(()), so the unknown-rank
case
+ # reports None and becomes a tensor with no static shape.
+ if not tensor_type.HasField("shape"):
+ return info_proto.name, None, elem_dtype, None, value_dict
+
shape = []
shape_name = []
for dim in info_proto.type.tensor_type.shape.dim:
@@ -1641,7 +1652,11 @@ class Shape(OnnxOpConverter):
return relax.ShapeExpr([data_info.ndim])
# If no shape is defined in the type, it must be computed at runtime.
- if not data_info.shape:
+ # A rank-0 tensor has a defined but empty shape, and an empty ShapeExpr
+ # is falsy, so compare against None instead of testing truthiness.
+ # Otherwise a scalar input takes the runtime path and downstream
+ # converters that match on relax.ShapeExpr see an opaque value.
+ if data_info.shape is None:
data_shape = bb.normalize(relax.op.shape_of(inputs[0]))
return data_shape
@@ -5954,6 +5969,11 @@ class ONNXGraphImporter:
self._input_names.append(i_name)
if i_name in self._shape:
i_shape = self._shape[i_name]
+ elif i_shape is None:
+ warnings.warn(
+ f"Input {i_name} has unknown rank. "
+ "Specifying a static shape may improve performance"
+ )
else:
if "?" in str(i_shape):
warning_msg = (
diff --git a/tests/python/relax/test_frontend_onnx.py
b/tests/python/relax/test_frontend_onnx.py
index a2b7bfb1cd..439c4374c5 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -3067,6 +3067,103 @@ def test_shape():
tvm.ir.assert_structural_equal(tvm_model, Expected)
+def test_shape_scalar_input():
+ # A rank-0 input has a known, empty static shape. It used to be imported as
+ # a runtime R.shape_of because an empty ShapeExpr is falsy, which made the
+ # result opaque to every converter that matches on relax.ShapeExpr.
+ shape_node = helper.make_node("Shape", ["data"], ["output"])
+
+ graph = helper.make_graph(
+ [shape_node],
+ "shape_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_scalar_test")
+ tvm_model = from_onnx(model, keep_params_in_input=True)
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(data: R.Tensor((), dtype="float32")) -> R.Shape([]):
+ R.func_attr({"num_input": 1})
+ with R.dataflow():
+ gv: R.Shape([]) = R.shape([])
+ R.output(gv)
+ return gv
+
+ tvm.ir.assert_structural_equal(tvm_model, Expected)
+
+
+def test_shape_unknown_rank_input():
+ # An input whose ValueInfoProto carries no shape field has unknown rank,
which
+ # must stay distinct from a rank-0 tensor. It has no static shape to fold,
so
+ # Shape has to keep the runtime path rather than reporting R.shape([]).
+ shape_node = helper.make_node("Shape", ["data"], ["output"])
+
+ data_vi = helper.make_tensor_value_info("data", TensorProto.FLOAT, None)
+ assert not data_vi.type.tensor_type.HasField("shape"), "test needs an
absent shape field"
+
+ graph = helper.make_graph(
+ [shape_node],
+ "shape_unknown_rank_test",
+ inputs=[data_vi],
+ outputs=[helper.make_tensor_value_info("output", TensorProto.INT64,
None)],
+ )
+
+ model = helper.make_model(graph, producer_name="shape_unknown_rank_test")
+ tvm_model = from_onnx(model, keep_params_in_input=True)
+
+ # The input keeps an unknown shape rather than collapsing to R.Tensor(()).
+ data_ty = tvm_model["main"].params[0].ty
+ assert data_ty.shape is None
+ assert data_ty.ndim == -1
+
+ # And Shape falls back to computing it at runtime.
+ 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 "relax.shape_of" in op_names
+
+
+def test_slice_of_scalar_shape():
+ # Slice consuming Shape of a rank-0 input used to raise "Slice requires a
+ # statically known input rank", because Shape handed it an opaque value
+ # instead of a ShapeExpr. ONNX Runtime returns an empty int64 tensor here.
+ nodes = [
+ helper.make_node("Shape", ["data"], ["shape"]),
+ helper.make_node("Slice", ["shape", "starts", "ends"], ["output"]),
+ ]
+
+ graph = helper.make_graph(
+ nodes,
+ "slice_of_scalar_shape_test",
+ inputs=[
+ helper.make_tensor_value_info("data", TensorProto.FLOAT, []),
+ ],
+ outputs=[helper.make_tensor_value_info("output", TensorProto.INT64,
[0])],
+ initializer=[
+ helper.make_tensor("starts", TensorProto.INT64, [1], [0]),
+ helper.make_tensor("ends", TensorProto.INT64, [1], [1]),
+ ],
+ )
+
+ model = helper.make_model(graph,
producer_name="slice_of_scalar_shape_test")
+ tvm_model = from_onnx(model)
+
+ output_ty = tvm_model["main"].ret_ty
+ assert isinstance(output_ty, relax.TensorType)
+ assert [int(dim) for dim in output_ty.shape] == [0]
+ assert output_ty.dtype == "int64"
+
+
@pytest.mark.parametrize(
"attrs,expected_shape",
[
@@ -3221,6 +3318,10 @@ def test_shape_start_end_scalar():
assert relax.analysis.check_well_formed(tvm_model)
+ # A rank-0 input has a known, empty static shape, so start=1 slices an
empty
+ # ShapeExpr and folds at import time. This used to fall back to a runtime
+ # shape_of / shape_to_tensor / strided_slice / tensor_to_shape chain,
because
+ # the empty ShapeExpr tested as falsy in Shape._impl_v13.
op_names = []
def collect_ops(expr):
@@ -3229,12 +3330,19 @@ def test_shape_start_end_scalar():
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",
- ]
+ assert op_names == []
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(data: R.Tensor((), dtype="float32")) -> R.Shape([]):
+ R.func_attr({"num_input": 1})
+ with R.dataflow():
+ gv: R.Shape([]) = R.shape([])
+ R.output(gv)
+ return gv
+
+ tvm.ir.assert_structural_equal(tvm_model, Expected)
def test_trilu():
