adityasingh2400 opened a new pull request, #20092:
URL: https://github.com/apache/tvm/pull/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
--
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]
