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 5cf92b27c6 [Fix][Relax][ONNX] Where: broadcast size-1 shape
expressions, materialize ShapeExpr inputs (#20210)
5cf92b27c6 is described below
commit 5cf92b27c68fe7c55b826a81900e387c8754fc3a
Author: HuEnwei <[email protected]>
AuthorDate: Sat Aug 29 12:04:23 2026 +0800
[Fix][Relax][ONNX] Where: broadcast size-1 shape expressions, materialize
ShapeExpr inputs (#20210)
Fixes: #20186
## Summary
The Relax ONNX frontend's `Where` importer (`Where._impl_v16`)
mishandled two
legal models that feed a **shape tensor** (the output of a `Shape` op,
imported
as a `relax.ShapeExpr`) into `Where`:
1. **Size-1 broadcasting in the all-constant/shape-expression path is
rejected.**
The fast path required the `condition`, `x` and `y` element lists to
have
exactly equal lengths, but ONNX `Where` follows NumPy-style
multi-directional
broadcasting, in which a size-1 dimension broadcasts against any length.
Models like
`Where(cond=[True], x=Shape((2,3))=[2,3], y=[4,5])` (legal in
onnxruntime)
failed to import with:
```
ValueError: Cannot broadcast condition to x and y
```
2. **A `ShapeExpr` mixed with a runtime tensor crashes the general
path.**
When at least one input is a runtime tensor (e.g. `condition` is a graph
input) and another input is a shape expression, the `ShapeExpr` was
passed
straight to `relax.op.where`, which only accepts tensors:
```
InternalError: Op(relax.where) requires argument 1 (x1) to be a tensor
```
## Root cause
`Where._impl_v16` has three branches. The shape-like branch (all inputs
are
`relax.Constant` or `relax.ShapeExpr`) compared element-list lengths for
exact
equality instead of applying NumPy-style broadcasting. The general
fallback
(`relax.op.where(inputs[0], inputs[1], inputs[2])`) never converted a
`relax.ShapeExpr` input into a tensor, so a shape expression
co-occurring with
a graph input reached `relax.op.where` as a non-tensor expression.
## Fix
In `Where._impl_v16`:
- **Shape-like path**: broadcast the 1-D element lists to a common
length,
repeating a length-1 list to that length (mirroring NumPy's size-1
broadcasting), then select elementwise. The output stays a shape
expression so
that downstream shape consumers (e.g. a `Reshape`'s shape input) keep
working.
If `get_prim_expr_list` raises (e.g. a rank-2 constant mixed with a
shape
tensor), the code falls through to the tensor path instead of aborting
import.
- **General path**: materialize any `relax.ShapeExpr` input into an
int64
tensor with `relax.op.shape_to_tensor` before calling `relax.op.where`,
so the
operator always operates on tensors.
## Tests
Added to `tests/python/relax/test_frontend_onnx.py`:
- `test_where_shape_expr` — correctness (via `check_correctness` against
onnxruntime) for a `Where` fed by a `Shape` output, both with
`condition` as a
size-1 initializer and as a graph input (regression for defect 2).
- `test_where_shape_expr_broadcast` — the all-constant/shape-expression
fast
path must broadcast size-1 dims (parametrized over `(1,)` and
equal-length
`condition`/`y` combinations; regression for defect 1).
## Validation
Differential testing against onnxruntime (over the `Shape`-tensor
`Where`
cases: size-1 broadcast, equal lengths, mixed graph-input `condition`,
rank-2
constant fallback) all match; no regressions on the already-supported
cases.
All 6 new/related `Where` tests pass in-tree, and the full
`test_frontend_onnx.py` suite shows no new failures (the pre-existing
`test_topk`/`test_unique` failures in the local environment are
unrelated to
`Where`).
---
python/tvm/relax/frontend/onnx/onnx_frontend.py | 34 +++++++--
tests/python/relax/test_frontend_onnx.py | 92 +++++++++++++++++++++++++
2 files changed, 120 insertions(+), 6 deletions(-)
diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py
b/python/tvm/relax/frontend/onnx/onnx_frontend.py
index 0bbfb4e281..0b2e5063c8 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ -1678,13 +1678,35 @@ class Where(OnnxOpConverter):
np_inputs = [inp.data.numpy() for inp in inputs]
output = _np.where(*np_inputs)
return relax.const(output, output.dtype)
+
if all([isinstance(inp, relax.Constant | relax.ShapeExpr) for inp in
inputs]):
- condition, x, y = [get_prim_expr_list(inp) for inp in inputs]
- if len(condition) != len(x) or len(condition) != len(y):
- raise ValueError("Cannot broadcast condition to x and y")
- output = [x if c else y for c, x, y in zip(condition, x, y)]
- return relax.ShapeExpr(output)
- return relax.op.where(inputs[0], inputs[1], inputs[2])
+ try:
+ condition, x, y = [get_prim_expr_list(inp) for inp in inputs]
+ except ValueError:
+ # Not a 1-D shape-like input (e.g. a rank-2 constant mixed with
+ # a shape tensor): fall through to the tensor path below.
+ condition = x = y = None
+ else:
+ n = max(len(condition), len(x), len(y))
+ if not all(len(v) in (1, n) for v in (condition, x, y)):
+ raise ValueError(
+ "Cannot broadcast condition, x and y with lengths "
+ f"{len(condition)}, {len(x)}, {len(y)}"
+ )
+ if n > 1:
+ condition = condition * n if len(condition) == 1 else
condition
+ x = x * n if len(x) == 1 else x
+ y = y * n if len(y) == 1 else y
+ output = [x if c else y for c, x, y in zip(condition, x, y)]
+ return relax.ShapeExpr(output)
+
+ tensors = []
+ for inp in inputs:
+ if isinstance(inp, relax.ShapeExpr):
+ tensors.append(relax.op.shape_to_tensor(inp))
+ else:
+ tensors.append(inp)
+ return relax.op.where(tensors[0], tensors[1], tensors[2])
class Clip(OnnxOpConverter):
diff --git a/tests/python/relax/test_frontend_onnx.py
b/tests/python/relax/test_frontend_onnx.py
index 984e05aff9..00c34b775b 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -3380,6 +3380,98 @@ def test_equal():
check_correctness(model)
+def _make_where_shape_expr_model(t1_shape, cond_vals, y_vals,
cond_is_input=False):
+ """Build an ONNX model whose Where input X comes from a Shape op.
+
+ X = Shape(t1) yields the int64 shape tensor of t1 (e.g. [2, 3]); cond and Y
+ are length-1 or length-rank initializers (or cond is a graph input). Where
+ must follow NumPy-style broadcasting over the element lists.
+ """
+ inits, gin, nodes = [], [], []
+ gin.append(helper.make_tensor_value_info("__t", TensorProto.FLOAT,
list(t1_shape)))
+ nodes.append(helper.make_node("Shape", ["__t"], ["X"]))
+ if cond_is_input:
+ gin.append(helper.make_tensor_value_info("cond", TensorProto.BOOL,
[len(cond_vals)]))
+ else:
+ inits.append(
+ helper.make_tensor(
+ "cond",
+ TensorProto.BOOL,
+ [len(cond_vals)],
+ list(np.array(cond_vals, dtype=np.bool_)),
+ )
+ )
+ inits.append(
+ helper.make_tensor(
+ "Y", TensorProto.INT64, [len(y_vals)], list(np.array(y_vals,
dtype=np.int64))
+ )
+ )
+ nodes.append(helper.make_node("Where", ["cond", "X", "Y"], ["Yout"]))
+ graph = helper.make_graph(
+ nodes,
+ "where_shape_expr_test",
+ gin,
+ [helper.make_tensor_value_info("Yout", TensorProto.INT64,
[len(t1_shape)])],
+ inits,
+ )
+ model = helper.make_model(graph, opset_imports=[helper.make_opsetid("",
16)])
+ model.ir_version = 8
+ return model
+
+
+def test_where_shape_expr():
+ """Where with a shape tensor (Shape op output) must import and run.
+
+ Regression: a shape expression mixed with a runtime tensor was passed
+ straight to relax.op.where, which raised 'Op(relax.where) requires argument
+ 1 (x1) to be a tensor'. The shape expression must be materialized to an
+ int64 tensor first.
+ """
+ # cond is an initializer (1,) -> broadcasts to (2,); x = Shape((2, 3)) =
[2, 3]
+ model = _make_where_shape_expr_model((2, 3), [True], [4, 5])
+ check_correctness(model, inputs={"__t": np.ones((2, 3), dtype="float32")},
opset=16)
+
+ # cond is a graph input (2,); x = Shape((2, 3)); y initializer (2,)
+ model = _make_where_shape_expr_model((2, 3), [True, False], [4, 5],
cond_is_input=True)
+ check_correctness(
+ model,
+ inputs={
+ "__t": np.ones((2, 3), dtype="float32"),
+ "cond": np.array([True, False], dtype=np.bool_),
+ },
+ opset=16,
+ )
+
+
[email protected](
+ "cond_vals, y_vals, expected",
+ [
+ ([True], [4, 5], [2, 3]), # cond (1,) broadcasts to (2,)
+ ([True, False], [4], [2, 4]), # y (1,) broadcasts to (2,)
+ ([True], [4], [2, 3]), # both (1,) broadcast to (2,)
+ ([True, False], [4, 5], [2, 5]), # equal lengths (regression)
+ ],
+)
+def test_where_shape_expr_broadcast(cond_vals, y_vals, expected):
+ """The all-constant/shape-expression fast path must broadcast size-1 dims.
+
+ Regression: length-1 element lists were rejected with 'Cannot broadcast
+ condition to x and y' instead of being broadcast to the common length, so
+ e.g. cond=[True] with x=[2, 3] failed to import.
+ """
+ t1_shape = (2, 3)
+ model = _make_where_shape_expr_model(t1_shape, cond_vals, y_vals)
+ feeds = {"__t": np.ones(t1_shape, dtype="float32")}
+ mod = from_onnx(model, shape_dict={"__t": list(t1_shape)})
+ with tvm.transform.PassContext(opt_level=3):
+ ex = tvm.compile(mod, target="llvm")
+ vm = relax.VirtualMachine(ex, tvm.cpu())
+ out = vm["main"](feeds["__t"])
+ np.testing.assert_array_equal(
+ np.array([int(i) for i in out], dtype=np.int64), np.array(expected,
dtype=np.int64)
+ )
+
+
def test_shape():
shape_node = helper.make_node("Shape", ["data"], ["output"])
