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 b16cdecbd8 [Fix][Relax][Frontend][ONNX] Validate Flatten axis range in
`from_onnx` (#20145)
b16cdecbd8 is described below
commit b16cdecbd84b8405d77069c50f0233e909eafc00
Author: HuEnwei <[email protected]>
AuthorDate: Sun Aug 30 12:56:08 2026 +0800
[Fix][Relax][Frontend][ONNX] Validate Flatten axis range in `from_onnx`
(#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]`.
Co-authored-by: FFChopon <[email protected]>
Co-authored-by: Claude <[email protected]>
---
python/tvm/relax/frontend/onnx/onnx_frontend.py | 13 +++++++++++++
1 file changed, 13 insertions(+)
diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py
b/python/tvm/relax/frontend/onnx/onnx_frontend.py
index 6e0d38a7f0..e8bd5eb0ae 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ -4397,6 +4397,19 @@ class Flatten(OnnxOpConverter):
def _impl_v13(cls, bb, inputs, attr, params):
axis = attr.get("axis", 1)
data_shape = list(inputs[0].ty.shape)
+ 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}"
+ )
if axis == 0:
new_shape = (1, -1)
