Use assert_close

Author: jackzhxng

extension/llm/modules/test/test_attention.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
 
@@ -94,8 +95,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}",
         )
 
@@ -127,7 +129,12 @@ def test_attention_eager(self):
         tt_res = self.tt_mha(
             self.x, self.x, input_pos=next_input_pos
         )  # Self attention with input pos.
-        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.
@@ -139,8 +146,10 @@ def test_attention_export(self):
         )
         et_res = et_mha_ep.module()(self.x, self.x, input_pos=self.input_pos)
         tt_res = self.tt_mha(self.x, self.x, input_pos=self.input_pos)
-        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}",
         )
 
@@ -168,8 +177,9 @@ def test_attention_executorch(self):
         et_res = method.execute((self.x, self.x, self.input_pos))
         tt_res = self.tt_mha(self.x, self.x, input_pos=self.input_pos)
 
-        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]}",
         )