siyiweigeHEW opened a new pull request, #20145:
URL: https://github.com/apache/tvm/pull/20145
Fixes: #20144
## Summary
The Relax ONNX frontend silently accepted a `Flatten` node whose `axis`
attribute is outside `[-r, r]` (where `r` is the input rank), producing a
wrong output shape. The ONNX spec requires `axis ∈ [-r, r]`, and onnxruntime
rejects such models with a `ShapeInferenceError`. This PR makes the frontend
reject out-of-range `axis` like onnxruntime.
## Root cause
`Flatten._impl_v13` in `python/tvm/relax/frontend/onnx/onnx_frontend.py`
computed the batch size as `data_shape[0:axis]` with no range check. Because
Python slicing silently clamps out-of-range indices, `axis=5` on a rank-3
input `(2, 3, 4)` sliced the whole shape, giving `(24, -1)` → output `(24,
1)`,
instead of raising an error.
## Fix
Normalize negative `axis` (`axis += rank`) and raise `ValueError` when the
result is outside `[0, rank]`:
```python
rank = len(data_shape)
# ONNX Flatten spec: "The value for axis must be in the range [-r, r], where
r
# is the rank of the input tensor. Negative value means counting dimensions
from
# the back." Normalize negative axis and validate the range, matching
onnxruntime
# which rejects out-of-range axis with a ShapeInferenceError.
if axis < 0:
axis += rank
if not 0 <= axis <= rank:
raise ValueError(
f"Flatten axis {attr.get('axis', 1)} is out of range [-{rank},
{rank}] "
f"for an input of rank {rank}"
)
```
## Validation
Differential test: Relax (build + `VirtualMachine`) vs onnxruntime.
| Case | onnxruntime | TVM before | TVM after | Result |
|---|---|---|---|---|
| `axis=5` on `(2,3,4)` | rejects (ShapeInferenceError) | silently `(24, 1)`
| raises `ValueError` | fixed |
| `axis=-4` on `(2,3,4)` | rejects (ShapeInferenceError) | silently `(1,
24)` | raises `ValueError` | fixed |
Regression (all valid `axis ∈ [-r, r]` unchanged, 0 differences vs
onnxruntime):
- 125 static cases (17 shapes × all valid axes), all pass `onnx.checker`
- 12 multi-dtype cases (float32 / int64 / bool)
- 9 dynamic-symbolic cases (`['N',3,4,5]`, axis `-4..4`)
Run:
```bash
python results/TVM/deepseek-v4-flash/prove_hum/onnx_Flatten/4严格_穷举差分.py
python results/TVM/deepseek-v4-flash/prove_hum/onnx_Flatten/minimal_repro.py
# expect ValueError now
```
## Files changed
- `python/tvm/relax/frontend/onnx/onnx_frontend.py` — `Flatten._impl_v13`:
normalize negative `axis` and raise `ValueError` for `axis ∉ [-r, r]`.
--
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]
