[Bug] QDQ Add and Mul fusion rejects valid per-axis quantization graphs #33335

import numpy as np
import onnx
import onnxruntime as ort
from onnx import helper, numpy_helper

def run_case(name, model, feeds, expected, disabled_pass):
    onnx.checker.check_model(model, full_check=True)
    print(name)
    def show(outputs):
        return [(a.shape, str(a.dtype), a.reshape(-1)[:12].tolist()) for a in outputs]
    print("expected:", show(expected))
    for mode in ("disable", "all", "ablate"):
        options = ort.SessionOptions()
        options.log_severity_level = 4
        options.intra_op_num_threads = options.inter_op_num_threads = 1
        options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL \
            if mode == "disable" else ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        try:
            session = ort.InferenceSession(model.SerializeToString(), options,
                providers=["CPUExecutionProvider"],
                disabled_optimizers={disabled_pass} if mode == "ablate" else set())
            actual = session.run(None, feeds)
            matches = all(a.shape == b.shape and a.dtype == b.dtype and
                          np.array_equal(a, b, equal_nan=True)
                          for a, b in zip(actual, expected))
            print(mode, show(actual), "matches expected:", matches)
        except Exception as error:
            print(mode, str(error))

model = helper.make_model(helper.make_graph(
    [
        helper.make_node('DequantizeLinear', ['Aq', 'a_s', 'a_zp'], ['Af'], axis=0),
        helper.make_node('DequantizeLinear', ['Bq', 'b_s', 'b_zp'], ['Bf']),
        helper.make_node('Add', ['Af', 'Bf'], ['Yf']),
        helper.make_node('QuantizeLinear', ['Yf', 'y_s', 'y_zp'], ['Y'])
    ], "repro",
    [helper.make_tensor_value_info('Aq', 2, [2, 3])],
    [helper.make_tensor_value_info('Y', 2, [2, 3])],
    [
        numpy_helper.from_array(np.array([0.10000000149011612, 0.20000000298023224], dtype=np.float32), 'a_s'),
        numpy_helper.from_array(np.array([1, 2], dtype=np.uint8), 'a_zp'),
        numpy_helper.from_array(np.array(0.05000000074505806, dtype=np.float32), 'y_s'),
        numpy_helper.from_array(np.array(128, dtype=np.uint8), 'y_zp'),
        numpy_helper.from_array(np.array([[10, 20, 30], [40, 50, 60]], dtype=np.uint8), 'Bq'),
        numpy_helper.from_array(np.array(0.05000000074505806, dtype=np.float32), 'b_s'),
        numpy_helper.from_array(np.array(0, dtype=np.uint8), 'b_zp')
    ]),
    opset_imports=[helper.make_opsetid('', 18)], ir_version=10)
run_case('Add', model, {'Aq': np.array([[1, 2, 30], [4, 5, 60]], dtype=np.uint8)},
    [np.array([[138, 150, 216], [176, 190, 255]], dtype=np.uint8)], 'QDQSelectorActionTransformer')

model = helper.make_model(helper.make_graph(
    [
        helper.make_node('DequantizeLinear', ['Aq', 'a_s', 'a_zp'], ['Af'], axis=0),
        helper.make_node('DequantizeLinear', ['Bq', 'b_s', 'b_zp'], ['Bf']),
        helper.make_node('Mul', ['Af', 'Bf'], ['Yf']),
        helper.make_node('QuantizeLinear', ['Yf', 'y_s', 'y_zp'], ['Y'])
    ], "repro",
    [helper.make_tensor_value_info('Aq', 2, [2, 3])],
    [helper.make_tensor_value_info('Y', 2, [2, 3])],
    [
        numpy_helper.from_array(np.array([0.10000000149011612, 0.20000000298023224], dtype=np.float32), 'a_s'),
        numpy_helper.from_array(np.array([1, 2], dtype=np.uint8), 'a_zp'),
        numpy_helper.from_array(np.array(0.05000000074505806, dtype=np.float32), 'y_s'),
        numpy_helper.from_array(np.array(128, dtype=np.uint8), 'y_zp'),
        numpy_helper.from_array(np.array([[10, 20, 30], [40, 50, 60]], dtype=np.uint8), 'Bq'),
        numpy_helper.from_array(np.array(0.05000000074505806, dtype=np.float32), 'b_s'),
        numpy_helper.from_array(np.array(0, dtype=np.uint8), 'b_zp')
    ]),
    opset_imports=[helper.make_opsetid('', 18)], ir_version=10)
run_case('Mul', model, {'Aq': np.array([[1, 2, 30], [4, 5, 60]], dtype=np.uint8)},
    [np.array([[128, 130, 215], [144, 158, 255]], dtype=np.uint8)], 'QDQSelectorActionTransformer')