The GitHub Actions job "CI" on tvm.git/main has succeeded.
Run started by GitHub user tlopex (triggered by tlopex).

Head commit for run:
5cf92b27c68fe7c55b826a81900e387c8754fc3a / HuEnwei <[email protected]>
[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`).

Report URL: https://github.com/apache/tvm/actions/runs/33232933215

With regards,
GitHub Actions via GitBox


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to