Use assert_close

Author: jackzhxng

extension/llm/modules/test/test_mha.py

@@ -13,6 +13,7 @@
     MultiHeadAttention as ETMultiHeadAttention,
 )
 from executorch.runtime import Runtime
+from torch.testing import assert_close
 from torchtune.models.llama3_1._position_embeddings import Llama3ScaledRoPE
 from torchtune.modules.attention import MultiHeadAttention as TTMultiHeadAttention
 
@@ -92,8 +93,9 @@ def test_attention_eager(self):
         et_res = self.et_mha(self.x, self.x)  # Self attention.
         tt_res = self.tt_mha(self.x, self.x)  # Self attention.
 
-        self.assertTrue(
-            torch.allclose(et_res, tt_res),
+        assert_close(
+            et_res,
+            tt_res,
             msg=f"TorchTune output is not close to ET output.\n\nTorchTune: {tt_res}\nET output: {et_res}",
         )
 
@@ -104,7 +106,11 @@ def test_attention_eager(self):
         # et_res = self.et_mha(self.x, self.x) # Self attention.
         # tt_res = self.tt_mha(self.x, self.x) # Self attention.
 
-        # self.assertTrue(torch.allclose(et_res, tt_res))
+        # assert_close(
+        #     et_res,
+        #     tt_res,
+        #     msg=f"TorchTune output is not close to ET output.\n\nTorchTune: {tt_res}\nET output: {et_res}"
+        # )
 
     def test_attention_export(self):
         # Self attention.
@@ -116,8 +122,10 @@ def test_attention_export(self):
         )
         et_res = et_mha_ep.module()(self.x, self.x)
         tt_res = self.tt_mha(self.x, self.x)
-        self.assertTrue(
-            torch.allclose(et_res, tt_res),
+
+        assert_close(
+            et_res,
+            tt_res,
             msg=f"TorchTune output is not close to ET output.\n\nTorchTune: {tt_res}\nET output: {et_res}",
         )
 
@@ -145,8 +153,9 @@ def test_attention_executorch(self):
         et_res = method.execute((self.x, self.x))
         tt_res = self.tt_mha(self.x, self.x)
 
-        self.assertTrue(
-            torch.allclose(et_res[0], tt_res, atol=1e-05),
+        assert_close(
+            et_res[0],
+            tt_res,
             msg=f"TorchTune output is not close to ET output.\n\nTorchTune: {tt_res}\nET output: {et_res[0]}",
         )