diff --git a/tools/Polygraphy/CHANGELOG.md b/tools/Polygraphy/CHANGELOG.md index cea3567ee8..4c1986b65c 100644 --- a/tools/Polygraphy/CHANGELOG.md +++ b/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 diff --git a/tools/Polygraphy/polygraphy/comparator/compare.py b/tools/Polygraphy/polygraphy/comparator/compare.py index c1376f2b77..7513b5d489 100644 --- a/tools/Polygraphy/polygraphy/comparator/compare.py +++ b/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)) diff --git a/tools/Polygraphy/tests/comparator/test_compare.py b/tools/Polygraphy/tests/comparator/test_compare.py index 31bb9dec30..09f65ab862 100644 --- a/tools/Polygraphy/tests/comparator/test_compare.py +++ b/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"]