Preserve Polygraphy cosine similarity under rescaling

tools/Polygraphy/CHANGELOG.md

@@ -2,6 +2,10 @@
 
 Dates are in YYYY-MM-DD format.
 
+## Unreleased
+### Fixed
+- Preserve cosine similarity under independent rescaling of finite output vectors, avoiding spurious comparison passes and failures.
+
 
 ## v0.53.6 (2026-09-22)
 ### Added

tools/Polygraphy/polygraphy/comparator/compare.py

@@ -1435,21 +1435,28 @@ def _compute_metric(self, out0, out1):
         array1_flat = util.array.ravel(comp_util.cast_up(out0))
         array2_flat = util.array.ravel(comp_util.cast_up(out1))
 
-        # Calculate dot product
-        dot_product = util.array.sum(util.array.multiply(array1_flat, array2_flat))
+        # Scale each vector before products to avoid overflow and underflow.
+        def normalize_scale(array):
+            if util.array.dtype(array).is_integral:
+                array = util.array.cast(array, DataType.FLOAT64)
+            scale = (
+                max(abs(util.array.min(array)), abs(util.array.max(array)))
+                if util.array.size(array)
+                else 0
+            )
+            return (array / scale if scale != 0 else array), scale
 
-        # Calculate magnitudes
+        array1_flat, scale1 = normalize_scale(array1_flat)
+        array2_flat, scale2 = normalize_scale(array2_flat)
+        if scale1 == 0 and scale2 == 0:
+            return 1.0
+        elif scale1 == 0 or scale2 == 0:
+            return 0.0
+
+        dot_product = util.array.sum(util.array.multiply(array1_flat, array2_flat))
         magnitude1 = util.array.sqrt(util.array.sum(util.array.power(array1_flat, 2)))
         magnitude2 = util.array.sqrt(util.array.sum(util.array.power(array2_flat, 2)))
 
-        # Avoid division by zero
-        if magnitude1 == 0 and magnitude2 == 0:
-            return (
-                1.0  # If both vectors are zero, they are identical (similarity = 1.0)
-            )
-        elif magnitude1 == 0 or magnitude2 == 0:
-            return 0.0  # If only one vector is zero, they are orthogonal (similarity = 0.0)
-
         # Cosine similarity is dot_product / (magnitude1 * magnitude2)
         cosine_similarity = float(dot_product / (magnitude1 * magnitude2))
 

tools/Polygraphy/tests/comparator/test_compare.py

@@ -572,6 +572,61 @@ def test_per_output_tolerance(self):
 
 
 class TestCosineSimilarityCompareFunc:
+    @pytest.mark.parametrize(
+        "array_type", [np.array, build_torch], ids=["numpy", "torch"]
+    )
+    @pytest.mark.parametrize(
+        "dtype, scale0, scale1",
+        [
+            (np.float32, 1e-30, 1e-30),
+            (np.float32, 1e30, 1e30),
+            (np.float64, 1e-200, 1e200),
+            (np.float64, 1e200, 1e200),
+        ],
+    )
+    @pytest.mark.parametrize(
+        "other, expected", [([3.0, 4.0], 1.0), ([4.0, -3.0], 0.0), ([-3.0, -4.0], -1.0)]
+    )
+    def test_independent_rescaling(
+        self, array_type, dtype, scale0, scale1, other, expected
+    ):
+        lhs = array_type(np.array([3.0, 4.0], dtype=dtype) * scale0)
+        rhs = array_type(np.array(other, dtype=dtype) * scale1)
+        result = CosineSimilarityCompareFunc()(
+            IterationResult({"output": lhs}), IterationResult({"output": rhs})
+        )["output"]
+
+        assert np.isclose(result.cosine_similarity, expected, atol=1e-6)
+        assert bool(result) == (expected >= 0.997)
+
+    @pytest.mark.parametrize(
+        "array_type", [np.array, build_torch], ids=["numpy", "torch"]
+    )
+    @pytest.mark.parametrize(
+        "values0, values1, dtype, expected",
+        [
+            ([], [], np.float32, 1.0),
+            ([0.0, 0.0], [3e-30, 4e-30], np.float32, 0.0),
+            ([-2147483648, 0], [-2147483648, 0], np.int32, 1.0),
+            ([-2147483648, 0], [2147483647, 0], np.int32, -1.0),
+        ],
+    )
+    def test_scale_controls(self, array_type, values0, values1, dtype, expected):
+        lhs = array_type(values0, dtype=dtype)
+        rhs = array_type(values1, dtype=dtype)
+        original_lhs, original_rhs = (
+            util.array.to_numpy(lhs).copy(),
+            util.array.to_numpy(rhs).copy(),
+        )
+        result = CosineSimilarityCompareFunc()(
+            IterationResult({"output": lhs}), IterationResult({"output": rhs})
+        )["output"]
+
+        assert np.isclose(result.cosine_similarity, expected, atol=1e-6)
+        assert bool(result) == (expected >= 0.997)
+        np.testing.assert_array_equal(util.array.to_numpy(lhs), original_lhs)
+        np.testing.assert_array_equal(util.array.to_numpy(rhs), original_rhs)
+
     def test_identical_outputs(self):
         res0, res1 = _make_results([1.0, 2.0, 3.0], [1.0, 2.0, 3.0])
         result = CosineSimilarityCompareFunc()(res0, res1)["output"]