Register Flatten as Direct8Bit op in Python QDQ static quantizer

[Rishi-Dave]

Summary

• Register Flatten in QLinearOpsRegistry (Direct8BitOp) and QDQRegistry (QDQDirect8BitOp) so the Python QDQ quantizer no longer emits a redundant DequantizeLinear -> Flatten -> QuantizeLinear pair around layout-only Flatten nodes.
• Add test_op_flatten.py covering both QOperator and QDQ formats in u8/u8 and s8/s8.

Motivation

Fixes #21375. The C++ runtime optimizer was updated in #21376 to drop redundant Q/DQ pairs around Flatten, mirroring its existing Reshape handling. The Python-layer quantizer was never updated, so quantize_static / quantize_dynamic still emitted the redundant Q/DQ pair around Flatten in QDQ format. This PR aligns the Python tooling with the runtime behavior.

Flatten is a layout-only op with no arithmetic — its quantization semantics are identical to Reshape, Squeeze, and Unsqueeze, all of which are already registered as Direct8BitOp / QDQDirect8BitOp.

Changes

• onnxruntime/python/tools/quantization/registry.py: add "Flatten": Direct8BitOp to QLinearOpsRegistry and "Flatten": QDQDirect8BitOp to QDQRegistry.
• onnxruntime/test/python/quantization/test_op_flatten.py: new test file modeled on test_op_reshape.py covering u8/u8 and s8/s8 in both QOperator and QDQ formats; asserts no extra QuantizeLinear/DequantizeLinear is inserted around Flatten, and verifies numerical correctness via check_model_correctness.

Test Plan

• python -m pytest onnxruntime/test/python/quantization/test_op_flatten.py -v — 2 tests pass.
• python -m pytest onnxruntime/test/python/quantization/test_op_reshape.py -v — regression: 2 tests pass.
• python -m pytest onnxruntime/test/python/quantization/test_op_squeeze_unsqueeze.py -v — regression: 2 tests pass.
• lintrunner -a clean on changed files.

[Copilot]

Pull request overview

This PR updates the Python quantization toolchain so Flatten is treated as a layout-only “direct 8-bit” op (similar to Reshape/Squeeze/Unsqueeze), avoiding redundant DequantizeLinear -> Flatten -> QuantizeLinear patterns in quantized models. It also adds a new Python unit test to validate Flatten behavior in both QOperator and QDQ quantization formats.

Changes:

• Register Flatten as Direct8BitOp in QLinearOpsRegistry and as QDQDirect8BitOp in QDQRegistry.
• Add test_op_flatten.py to validate node patterns and numerical correctness for u8/u8 and s8/s8 in both QOperator and QDQ modes.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Description
onnxruntime/python/tools/quantization/registry.py	Adds Flatten to the direct-8bit registries for QOperator and QDQ quantization flows.
onnxruntime/test/python/quantization/test_op_flatten.py	Adds tests intended to ensure Flatten does not get redundant Q/DQ pairs and that quantized models are numerically correct.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Rishi-Dave]

A few notes on the Copilot review pass to pre-empt likely questions:

Scope (registry.py): The PR description was slightly loose — the change only aligns quantize_static / QDQ behavior. quantize_dynamic operates via IntegerOpsRegistry (QuantizationMode.IntegerOps), where Flatten is intentionally not registered: dynamic int quantization there is restricted to ops with real integer kernels (MatMulInteger, ConvInteger, etc.), and Flatten has no such kernel. The right fix is in QLinearOpsRegistry + QDQRegistry, matching how Reshape / Squeeze / Unsqueeze / Transpose are already registered as Direct8Bit. Happy to tighten the PR description if maintainers prefer.

Test shape (test_op_flatten.py): The model uses a 2D MatMul output with axis=1, which exercises the QDQ insertion path (Q/DQ around Flatten input/output, shared scale) regardless of whether shape changes numerically. I can extend it to a 3D MatMul + axis=1 (true rank reduction) in a follow-up if reviewers want stronger coverage.

CodeQL duplicate import: trivial — will consolidate to a single from onnx import ... if requested.

[Rishi-Dave]

Addressed the CodeQL py/import-and-import-from finding in onnxruntime/test/python/quantization/test_op_flatten.py (alert #34949).

Dropped from onnx import TensorProto, helper and routed all references through the existing import onnx (onnx.TensorProto, onnx.helper.make_tensor_value_info, etc.). Pure mechanical rename — no test logic touched. lintrunner -a is clean.

Commit: 5d4947b

[tianleiwu]

Summary

The core change is correct and well-targeted. Flatten is a layout-only op with no arithmetic, so its quantization semantics are identical to Reshape/Squeeze/Unsqueeze, and routing it through Direct8BitOp/QDQDirect8BitOp (which reuse input[0]'s scale/zero-point for the output) is the right behavior for any axis. Registering it in QLinearOpsRegistry and QDQRegistry correctly removes the redundant DequantizeLinear -> Flatten -> QuantizeLinear pair, aligning the Python QDQ tooling with the runtime optimizer from #21376. The new test closely mirrors test_op_reshape.py, and the check_op_nodes assertion (Flatten input rewired to the quantized tensor) plus the QDQ count assertions (QuantizeLinear: 3, DequantizeLinear: 4, matching the Reshape baseline) do verify that no extra Q/DQ is inserted.

Verdict: COMMENT — approve-worthy once the two minor items below are considered.

Suggestions (non-blocking)

1. Test exercises a shape-preserving Flatten. construct_model_matmul_flatten([3, 7], [7, 7], [3, 7]) with axis=1 produces output shape [3, 7] identical to the MatMul output, so Flatten is effectively an identity here. The redundant-Q/DQ-removal is still validated, but the test would be more representative of real exported models if it used a higher-rank MatMul output (e.g. a 3D batched MatMul) so Flatten(axis=1) actually collapses dimensions. (This matches the existing automated reviewer note on the construct_model_matmul_flatten call.)

2. PR description overclaims dynamic-quantization coverage. The description states the change aligns both quantize_static and quantize_dynamic, but only QLinearOpsRegistry/QDQRegistry are updated. quantize_dynamic uses QuantizationMode.IntegerOps (IntegerOpsRegistry), which is untouched, so dynamic quantization of Flatten is unchanged. Please either narrow the description to static/QDQ, or register Flatten for IntegerOps if dynamic coverage is intended.

For reference, the earlier CodeQL note about onnx being imported via both import and import from is already resolved on the current head (commit consolidated the onnx imports).