Clean up output validation helper

Summary: In helper `_assert_eager_lowered_same_result`, we want to allow user to validate different types of output from eager model and ET model. We can use a `validation_fn` to do this, and let user use the proper validation function to check.

Reviewed By: guangy10

Differential Revision: D48267539

fbshipit-source-id: abb32c0f5bd26ccefcecd3f8b56767c3ce5789f2

Author: kirklandsign

examples/export/test/test_export.py

@@ -6,6 +6,8 @@
 
 import unittest
 
+from typing import Any, Callable
+
 import torch
 
 from executorch.examples.export.utils import _CAPTURE_CONFIG, _EDGE_COMPILE_CONFIG
@@ -19,8 +21,17 @@
 
 class ExportTest(unittest.TestCase):
     def _assert_eager_lowered_same_result(
-        self, eager_model: torch.nn.Module, example_inputs
+        self,
+        eager_model: torch.nn.Module,
+        example_inputs,
+        validation_fn: Callable[[Any, Any], bool],
     ):
+        """
+        Asserts that the given model has the same result as the eager mode
+        lowered model, with example_inputs, validated by validation_fn, which
+        takes the eager mode output and ET output, and returns True if they
+        match.
+        """
         import executorch.exir as exir
 
         edge_model = exir.capture(eager_model, example_inputs, _CAPTURE_CONFIG).to_edge(
@@ -36,38 +47,54 @@ def _assert_eager_lowered_same_result(
         with torch.no_grad():
             executorch_output = pte_model.run_method("forward", example_inputs)
 
-        if isinstance(eager_output, tuple):
-            # TODO: Allow validating other items
-            self.assertTrue(
-                torch.allclose(
-                    eager_output[0], executorch_output[0][0], rtol=1e-5, atol=1e-5
-                )
-            )
-        else:
-            self.assertTrue(
-                torch.allclose(eager_output, executorch_output[0], rtol=1e-5, atol=1e-5)
-            )
+        self.assertTrue(validation_fn(eager_output, executorch_output))
+
+    @staticmethod
+    def validate_tensor_allclose(eager_output, executorch_output):
+        return torch.allclose(
+            eager_output,
+            executorch_output[0],
+            rtol=1e-5,
+            atol=1e-5,
+        )
 
     def test_mv3_export_to_executorch(self):
         eager_model, example_inputs = MODEL_NAME_TO_MODEL["mv3"]()
         eager_model = eager_model.eval()
 
-        self._assert_eager_lowered_same_result(eager_model, example_inputs)
+        self._assert_eager_lowered_same_result(
+            eager_model, example_inputs, self.validate_tensor_allclose
+        )
 
     def test_mv2_export_to_executorch(self):
         eager_model, example_inputs = MODEL_NAME_TO_MODEL["mv2"]()
         eager_model = eager_model.eval()
 
-        self._assert_eager_lowered_same_result(eager_model, example_inputs)
+        self._assert_eager_lowered_same_result(
+            eager_model, example_inputs, self.validate_tensor_allclose
+        )
 
     def test_emformer_export_to_executorch(self):
         eager_model, example_inputs = MODEL_NAME_TO_MODEL["emformer"]()
         eager_model = eager_model.eval()
 
-        self._assert_eager_lowered_same_result(eager_model, example_inputs)
+        validate_emformer_result = (
+            lambda eager_output, executorch_output: torch.allclose(
+                eager_output[0],
+                executorch_output[0][0],
+                rtol=1e-5,
+                atol=1e-5,
+            )
+        )
+
+        self._assert_eager_lowered_same_result(
+            eager_model, example_inputs, validate_emformer_result
+        )
 
     def test_vit_export_to_executorch(self):
         eager_model, example_inputs = MODEL_NAME_TO_MODEL["vit"]()
         eager_model = eager_model.eval()
 
-        self._assert_eager_lowered_same_result(eager_model, example_inputs)
+        self._assert_eager_lowered_same_result(
+            eager_model, example_inputs, self.validate_tensor_allclose
+        )