[Fix][Relax][Frontend][ONNX] Validate Flatten axis range in from_onnx

[siyiweigeHEW]

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]:

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:

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