Add MSE numerical comparator

Differential Revision: D76781331

Pull Request resolved: #11759

Author: Juntian777

devtools/inspector/numerical_comparator/TARGETS

@@ -9,7 +9,6 @@ python_library(
     deps = [],
 )
 
-
 python_library(
     name = "l1_numerical_comparator",
     srcs = ["l1_numerical_comparator.py"],
@@ -19,12 +18,20 @@ python_library(
     ],
 )
 
-
+python_library(
+    name = "mse_numerical_comparator",
+    srcs = ["mse_numerical_comparator.py"],
+    deps = [
+        "//executorch/devtools/inspector/numerical_comparator:numerical_comparator_base",
+        "//executorch/devtools/inspector:lib",
+    ],
+)
 
 python_library(
     name = "lib",
     srcs = ["__init__.py"],
     deps = [
         ":l1_numerical_comparator",
+        ":mse_numerical_comparator",
     ],
 )

devtools/inspector/numerical_comparator/__init__.py

@@ -9,5 +9,9 @@
     L1Comparator,
 )
 
+from executorch.devtools.inspector.numerical_comparator.mse_numerical_comparator import (
+    MSEComparator,
+)
+
 
-__all__ = ["L1Comparator"]
+__all__ = ["L1Comparator", "MSEComparator"]

devtools/inspector/numerical_comparator/mse_numerical_comparator.py

@@ -0,0 +1,32 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Any
+
+import torch
+from executorch.devtools.inspector._inspector_utils import convert_to_float_tensor
+from executorch.devtools.inspector.numerical_comparator.numerical_comparator_base import (
+    NumericalComparatorBase,
+)
+
+
+class MSEComparator(NumericalComparatorBase):
+    def compare(self, a: Any, b: Any) -> float:
+        """Compare mean squared difference between two outputs."""
+
+        t_a = convert_to_float_tensor(a)
+        t_b = convert_to_float_tensor(b)
+        if torch.isnan(t_a).any() or torch.isnan(t_b).any():
+            t_a = torch.nan_to_num(t_a)
+            t_b = torch.nan_to_num(t_b)
+
+        try:
+            res = float(torch.mean(torch.square(t_a - t_b)))
+        except Exception as e:
+            raise ValueError(
+                f"Error computing MSE difference between tensors: {str(e)}"
+            )
+        return res

devtools/inspector/tests/TARGETS

@@ -62,6 +62,14 @@ python_unittest(
     ],
 )
 
+python_unittest(
+    name = "mse_comparator_test",
+    srcs = ["mse_comparator_test.py"],
+    deps = [
+        "//executorch/devtools/inspector/numerical_comparator:lib",
+    ],
+)
+
 python_library(
     name = "inspector_test_utils",
     srcs = [

devtools/inspector/tests/mse_comparator_test.py

@@ -0,0 +1,56 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import unittest
+
+import torch
+
+from executorch.devtools.inspector.numerical_comparator import MSEComparator
+
+
+class TestMSEComparator(unittest.TestCase):
+    mse_comparator = MSEComparator()
+
+    def test_identical_tensors(self):
+        a = torch.tensor([[10, 4], [3, 4]])
+        b = torch.tensor([[10, 4], [3, 4]])
+        expected = 0.0
+        result = self.mse_comparator.compare(a, b)
+        self.assertAlmostEqual(result, expected)
+
+    def test_scalar(self):
+        a = 10
+        b = 2
+        expected = 64.0
+        result = self.mse_comparator.compare(a, b)
+        self.assertAlmostEqual(result, expected)
+
+    def test_with_nans_replaced_with_zero(self):
+        a = torch.tensor([3, 1, -3, float("nan")])
+        b = torch.tensor([float("nan"), 0, -3, 2])
+        expected = (9.0 + 1.0 + 0.0 + 4.0) / 4.0
+        result = self.mse_comparator.compare(a, b)
+        self.assertAlmostEqual(result, expected)
+
+    def test_shape_mismatch_raises_exception(self):
+        a = torch.tensor([0, 2, -1])
+        b = torch.tensor([1, 1, -3, 4])
+        with self.assertRaises(ValueError):
+            self.mse_comparator.compare(a, b)
+
+    def test_2D_tensors(self):
+        a = torch.tensor([[4, 9], [6, 4]])
+        b = torch.tensor([[1, 2], [3, 10]])
+        expected = (9.0 + 49.0 + 9.0 + 36.0) / 4.0
+        result = self.mse_comparator.compare(a, b)
+        self.assertAlmostEqual(result, expected)
+
+    def test_list_of_tensors(self):
+        a = [torch.tensor([2, 4]), torch.tensor([15, 2])]
+        b = [torch.tensor([1, 2]), torch.tensor([9, 5])]
+        expected = (1.0 + 4.0 + 36.0 + 9.0) / 4.0
+        result = self.mse_comparator.compare(a, b)
+        self.assertAlmostEqual(result, expected)