Emit per-channel axis for fused bias quantizers in ONNX QDQ export (#7314)

* Emit per-channel axis for fused bias quantizers in ONNX QDQ export

Signed-off-by: Shivakumar, Shobitha <shobitha@qti.qualcomm.com>

* Source per-channel export axis from quant_info instead of tensor_quantizer_params

Signed-off-by: Shivakumar, Shobitha <shobitha@qti.qualcomm.com>

---------

Signed-off-by: Shivakumar, Shobitha <shobitha@qti.qualcomm.com>

Author: ShobithaShivakumar

TrainingExtensions/onnx/src/python/aimet_onnx/_encoding.py

@@ -697,12 +697,14 @@ def from_quantizer(cls, qtzr: QcQuantizeOp) -> AffineEncoding | None:
         if not qtzr.enabled:
             return None
 
-        if qtzr.quant_info.usePerChannelMode and qtzr.tensor_quantizer_params:
-            channel_axis = qtzr.tensor_quantizer_params.channel_axis
+        if qtzr.quant_info.usePerChannelMode:
+            # Per-channel scale is a vector, so an axis must be emitted alongside it.
+            # Source the axis from quant_info -- the same fields _encoding_shape() uses
+            # to shape the scale vector -- rather than tensor_quantizer_params, whose
+            # channel_axis is unset for fused/derived int32 bias quantizers.
+            channel_axis = qtzr.quant_info.channelAxis
             block_size = qtzr.quant_info.blockSize or None
-            block_axis = (
-                None if block_size is None else qtzr.tensor_quantizer_params.block_axis
-            )
+            block_axis = None if block_size is None else qtzr.quant_info.blockAxis
         else:
             channel_axis = None
             block_size = None

TrainingExtensions/onnx/test/python/test_quantsim.py

@@ -4922,6 +4922,53 @@ def test_matmul_add_bias_quantizer(self, per_channel: bool, fuse_supergroups):
         # Making sure that the qdq graph is runnable
         qdq_output = quantized_model_session.run(None, dummy_input)
 
+    @pytest.mark.parametrize("tqp_channel_axis", [None, 0])
+    def test_per_channel_export_emits_axis_without_tqp_axis(self, tqp_channel_axis):
+        """
+        A per-channel quantizer (vector scale) must export an `axis`, sourced from
+        quant_info.channelAxis, regardless of tensor_quantizer_params.channel_axis --
+        which is unset for the fused/derived int32 bias quantizers that triggered this.
+        Otherwise the DequantizeLinear is invalid for the 1-D bias tensor.
+        """
+        from aimet_onnx.common import libquant_info
+        from aimet_onnx.qc_quantize_op import QcQuantizeOp, TensorQuantizerParams
+
+        num_channels = 8
+        tqp = TensorQuantizerParams([num_channels])
+        tqp.channel_axis = tqp_channel_axis
+        tqp.block_axis = None
+
+        qtzr = QcQuantizeOp(
+            quant_info=libquant_info.QcQuantizeInfo(),
+            tensor_quantizer_params=tqp,
+            op_mode=libpymo.TensorQuantizerOpMode.oneShotQuantizeDequantize,
+        )
+        qtzr.enabled = True
+        qtzr.bitwidth = 32
+        qtzr.use_symmetric_encodings = True
+
+        # Put the quantizer in per-channel mode (vector scale shaped by channelAxis).
+        qtzr.quant_info.usePerChannelMode = True
+        qtzr.quant_info.channelAxis = 0
+        qtzr.quant_info.blockAxis = 0
+        qtzr.quant_info.blockSize = 0
+        qtzr._tensor_quantizer = qtzr._build_tensor_quantizer()
+
+        encodings = []
+        for scale in np.linspace(1e-4, 1e-3, num_channels):
+            enc = libpymo.TfEncoding()
+            enc.bw = 32
+            enc.delta = float(scale)
+            enc.offset = -(2**31)
+            enc.min = float(scale) * -(2**31)
+            enc.max = float(scale) * (2**31 - 1)
+            encodings.append(enc)
+        qtzr.load_encodings(encodings)
+
+        exported = qtzr.export_encodings("2.0.0")
+        assert len(exported["y_scale"]) == num_channels
+        assert exported.get("axis") == 0
+
     @pytest.mark.parametrize(
         "per_channel, weight_encoding",
         [