[Bug][Relax][Frontend][ONNX] Flatten with out-of-range `axis` silently accepted by `from_onnx`; onnxruntime rejects the same model

[siyiweigeHEW]

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). 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

"""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