Fix version check

Signed-off-by: Ganesan Ramalingam <[email protected]>

Author: gramalingam

onnxscript/rewriter/rules/fusion/_gqa_test.py

@@ -81,7 +81,11 @@ def test_basic_gqa_fusion(self):
         self.assertGreater(count, 0, "GQA fusion should have occurred")
 
         # We can't yet test numerical equivalence because of a bug in the op spec/implementation.
-        if version.parse(onnx.__version__) >= version.parse("1.19.1"):
+        onnx_ver = version.parse(onnx.__version__)
+        if onnx_ver >= version.parse("1.19.1") and not (
+            onnx_ver.is_prerelease or onnx_ver.is_devrelease
+        ):
+            # Only official releases >= 1.19.1
             onnxscript.optimizer.remove_unused_nodes(model)
             rewritten_model_proto = ir.serde.serialize_model(model)
             onnxscript.rewriter.testing.assert_numerically_equal(

onnxscript/rewriter/testing.py

@@ -102,7 +102,12 @@ def assert_numerically_equal(
     for i, (orig, rewritten) in enumerate(zip(original_outputs, the_rewritten_outputs)):
         print(f"==== Output {i} ====")
         diff = np.abs(orig - rewritten)
-        print(diff)
+        for h in range(diff.shape[1]):
+            subarray = diff[:, h, :, :]  # Select along H
+            if np.allclose(subarray, 0):
+                print(f"H={h}: all zeros")
+            else:
+                print(f"H={h}: not all zeros")
 
     np.testing.assert_allclose(
         original_outputs, the_rewritten_outputs, rtol=rtol, atol=atol, equal_nan=True