siyiweigeHEW opened a new issue, #20144: URL: https://github.com/apache/tvm/issues/20144
### Expected behavior The ONNX `Flatten` op requires the `axis` attribute to be in the range `[-r, r]`, where `r` is the rank of the input tensor (see [ONNX Flatten spec](https://github.com/onnx/onnx/blob/main/docs/Operators.md#Flatten)). A model whose `Flatten` node has `axis` outside this range is out-of-spec. The reference runtime onnxruntime rejects such models at session creation with a `ShapeInferenceError`: ``` [ONNXRuntimeError] : 1 : FAIL : Node () Op (Flatten) [ShapeInferenceError] Invalid value(5) for attribute 'axis' ``` `tvm.relax.frontend.onnx.from_onnx` should either reject out-of-range `axis` values (consistent with onnxruntime), or at minimum raise an error instead of silently computing a result. ### Actual behavior `from_onnx` silently accepts an invalid `Flatten` model with `axis=5` on a rank-3 input `X: (2, 3, 4)` and the built VM returns `(24, 1)` — a shape that only exists because the frontend slices `data_shape[0:5]` = the whole shape and feeds `(24, -1)` into `reshape`. No error or warning is raised. The same divergence happens for `axis=-4` on a rank-3 input (`(1, 24)`), and also when the input rank is symbolic (dynamic dims). The `Flatten` frontend implementation at `python/tvm/relax/frontend/onnx/onnx_frontend.py:2722` (`Flatten._impl_v13`) computes the batch size via `data_shape[0:axis]` without ever checking `|axis| <= r`. ### Environment - OS: Linux - TVM: v0.24.dev0 (commit `262c6d2e0`) - Python: 3.11 - onnx: 1.20.1 - onnxruntime: 1.24.1 ### Steps to reproduce ```python """Repro: ONNX Flatten with out-of-range axis (|axis| > rank) is silently accepted by TVM relax frontend, while onnxruntime rejects it with a ShapeInferenceError.""" import numpy as np import onnx, onnxruntime from onnx import helper, TensorProto import tvm from tvm import relax from tvm.relax.frontend.onnx import from_onnx # (2,3,4) input, axis=5 > rank(3): violates ONNX spec "axis must be in the range [-r, r]" X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [2, 3, 4]) Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [24, 1]) node = helper.make_node("Flatten", ["X"], ["Y"], axis=5) graph = helper.make_graph([node], "flat", [X], [Y]) model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)]) model.ir_version = 8 onnx.checker.check_model(model) # (1) checker passes print("onnx.checker: PASS") try: onnxruntime.InferenceSession(model.SerializeToString()) # (2) onnxruntime rejects print("onnxruntime: accepted (unexpected)") except Exception as e: print("onnxruntime rejects:", str(e).splitlines()[0]) mod = from_onnx(model, shape_dict={"X": [2, 3, 4]}) # (3) TVM accepts silently ex = relax.build(mod, target="llvm") vm = relax.VirtualMachine(ex, tvm.cpu()) y = vm["main"](np.arange(24, dtype="float32").reshape(2, 3, 4)).numpy() print("TVM relax frontend accepts -> output shape:", y.shape) ``` Actual output: ``` onnx.checker: PASS onnxruntime rejects: Fail [ONNXRuntimeError] : 1 : FAIL : Node () Op (Flatten) [ShapeInferenceError] Invalid value(5) for attribute 'axis' TVM relax frontend accepts -> output shape: (24, 1) ``` ### Additional context - The ONNX spec for `Flatten.axis` states: *"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."* - `onnx.checker` does not catch the invalid `axis` (it does not validate attribute ranges), so this is not caught at model-construction time either. - A simple guard in `Flatten._impl_v13` before slicing, e.g. checking `0 <= axis <= r` (after normalizing negatives as `axis = axis + r`), would bring TVM in line with onnxruntime. ### 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]
