From 83bfc841197147020cb7de813e0bfb04ed3b2573 Mon Sep 17 00:00:00 2001 From: FFChopon <3483776996@qq.com> Date: Tue, 18 Aug 2026 01:56:38 +0800 Subject: [PATCH] [Relax][Frontend][ONNX] Validate Flatten axis range in from_onnx (#20144) Flatten._impl_v13 now normalizes negative axis and rejects axis outside [-r, r] (the rank of the input tensor), matching onnxruntime's ShapeInferenceError. Previously an out-of-range axis (e.g. axis=5 on a rank-3 input) was silently accepted and produced a wrong output shape (24, 1) instead of being rejected. Fixes #20144 Co-Authored-By: Claude --- 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 6bbf220dfe2e..e57d496ae565 100644 --- a/python/tvm/relax/frontend/onnx/onnx_frontend.py +++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py @@ -4135,6 +4135,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)