siyiweigeHEW opened a new issue, #20186:
URL: https://github.com/apache/tvm/issues/20186
### Expected behavior
A valid ONNX `Where` node follows the ONNX spec's NumPy-style
multi-directional broadcasting for `condition`, `X`, and `Y`, including size-1
dimensions. In particular, when one of the inputs is a shape tensor (the output
of the `Shape` operator, represented by the Relax ONNX frontend as a
`relax.ShapeExpr`), valid models such as the two below must import and run —
they pass `onnx.checker` and run correctly in onnxruntime and `onnx.reference`:
- **Case 1** (`condition` is a length-1 constant that broadcasts): `t` of
shape `(2, 3)` → `Shape` → `X` = `(2,)` int64 tensor `[2, 3]`; `condition` =
const bool `(1,)` `[True]`; `Y` = const int64 `(2,)` `[4, 5]`. Output must be
`(2,)` → `[2, 3]`.
- **Case 2** (`condition` is a runtime tensor while `X` is a shape tensor):
`condition` = graph input bool `(2,)`; `X` = `Shape(t)` → `(2,)` `[2, 3]`; `Y`
= const int64 `(2,)` `[4, 5]`. Output must be `(2,)`.
### Actual behavior
`tvm.relax.frontend.onnx.from_onnx` rejects/crashes on both valid models:
1. **Size-1 broadcasting is rejected** (`Where._impl_v16`,
`python/tvm/relax/frontend/onnx/onnx_frontend.py:992-1003`). The shape-like
branch only fires when *all* inputs are `relax.Constant`/`relax.ShapeExpr`; it
requires `len(condition) == len(x) == len(y)` and otherwise raises:
```
ValueError: Cannot broadcast condition to x and y
```
No ONNX broadcasting is performed. A valid model (Case 1, condition has
one size-1 dim that broadcasts to length 2) that onnxruntime and
`onnx.reference` both accept and return `(2,)` `[2, 3]` fails to import.
2. **Mixed runtime-tensor condition + `ShapeExpr` input crashes** (same
converter, fall-through path). When the condition is a runtime graph input (a
`relax.Var`, not a constant), the converter skips the shape-like branch and
calls `relax.op.where(inputs[0], inputs[1], inputs[2])` with the `ShapeExpr`
still as a shape (never materialized as an int64 tensor), raising:
```
InternalError: Operator Op(relax.where) requires argument 1 (x1) to be a
tensor. However, the argument R.shape([2, 3]) is instead of type R.Shape([2,
3])
```
A valid model (Case 2) that onnxruntime and `onnx.reference` both return
`(2,)` `[2, 5]` crashes on import.
### Root cause
`Where._impl_v16` in
`python/tvm/relax/frontend/onnx/onnx_frontend.py:992-1003`:
```python
@classmethod
def _impl_v16(cls, bb, inputs, attr, params):
if all([isinstance(inp, relax.Constant) for inp in inputs]):
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])
```
- The `Shape` ONNX op returns a `relax.ShapeExpr` (a shape, not a tensor),
so any downstream *data*-use of a shape tensor reaches `Where` as a `ShapeExpr`.
- The shape-like branch implements element-wise selection over Python lists
and requires equal lengths — it has no notion of ONNX/NumPy broadcasting, so
valid size-1 broadcast models are rejected.
- The fall-through path passes a `ShapeExpr` directly to `relax.op.where`,
which requires tensor operands, crashing instead of converting the shape tensor
to an int64 tensor.
### Environment
- OS: Linux
- TVM: v0.24.dev0 (main branch, commit `262c6d2e0`, built 2026-02-11)
- Python: 3.11
- onnx: 1.20.1
- onnxruntime: 1.24.1
### Steps to reproduce
```python
"""Repro: ONNX Where with shape-tensor (ShapeExpr) inputs is mishandled by
the TVM relax frontend."""
import numpy as np
import onnx
from onnx import helper, TensorProto
from onnx.reference import ReferenceEvaluator
import onnxruntime
from tvm.relax.frontend.onnx import from_onnx
def make(t1_shape, cond_vals, y_vals, cond_is_input=False):
"""X = Shape(t1) (shape tensor), cond/Y constants (or cond a graph
input)."""
inits, gin, sh = [], [], []
gin.append(helper.make_tensor_value_info("__t", TensorProto.FLOAT,
list(t1_shape)))
sh.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))))
w = helper.make_node("Where", ["cond", "X", "Y"], ["Yout"])
Y = helper.make_tensor_value_info("Yout", TensorProto.INT64,
[len(t1_shape)])
g = helper.make_graph(sh + [w], "g", gin, [Y], inits)
m = helper.make_model(g, opset_imports=[helper.make_opsetid("", 16)])
m.ir_version = 8
return m
def show(label, m, feeds, shape_dict):
print("=" * 66)
print(label)
onnx.checker.check_model(m) # valid
ONNX model
r = onnxruntime.InferenceSession(m.SerializeToString(),
providers=["CPUExecutionProvider"]).run(None, feeds)[0]
print("onnxruntime :", r.shape, r.tolist())
rr = ReferenceEvaluator(m).run(None, feeds)[0]
print("onnx.reference:", rr.shape, rr.tolist())
try:
from_onnx(m, shape_dict=shape_dict)
print("TVM from_onnx : OK")
except Exception as e:
print(f"TVM from_onnx : {type(e).__name__}: {str(e)[:110]}")
# Case 1: condition const (1,)=[True] must broadcast to (2,); X =
Shape((2,3)) -> [2,3]
m1 = make((2, 3), [True], [4, 5])
show("Case 1: cond const (1,), X=Shape((2,3))->[2,3], Y const (2,)=[4,5]",
m1, {"__t": np.ones((2, 3), np.float32)}, {"__t": [2, 3]})
# Case 2: condition is a runtime graph input (2,); X = Shape((2,3)) -> [2,3]
m2 = make((2, 3), [True, False], [4, 5], cond_is_input=True)
show("Case 2: cond graph input (2,), X=Shape((2,3))->[2,3], Y const
(2,)=[4,5]",
m2, {"__t": np.ones((2, 3), np.float32), "cond": np.array([True,
False])},
{"__t": [2, 3], "cond": [2]})
```
Actual output:
```
==================================================================
Case 1: cond const (1,), X=Shape((2,3))->[2,3], Y const (2,)=[4,5]
onnxruntime : (2,) [2, 3]
onnx.reference: (2,) [2, 3]
Error converting operator Where, with inputs:
[metadata["relax.expr.Constant"][0] ...
TVM from_onnx : ValueError: Cannot broadcast condition to x and y
==================================================================
Case 2: cond graph input (2,), X=Shape((2,3))->[2,3], Y const (2,)=[4,5]
onnxruntime : (2,) [2, 5]
onnx.reference: (2,) [2, 5]
TVM from_onnx : InternalError: Operator Op(relax.where) requires argument 1
(x1) to be a tensor. However, the argument R.shape([2, 3]) is instead of type
R.Shape([2, 3])
```
### Scope note
The common `Where` usages (three ordinary tensors, or all-constant inputs,
including multi-directional broadcasting) are handled correctly — verified
against onnxruntime and `onnx.reference` (12/12 cases, `max|diff| = 0`). The
defect is limited to the shape-tensor (`ShapeExpr`) path, which requires a
`Where` input produced by the `Shape` operator (or another shape-computation
chain).
### Triage
* needs-triage
* bug
* relax
* frontend/onnx
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
