executorch/extension/llm/modules/test

Reviewed By: avikchaudhuri Differential Revision: D67383699
pytorch · Dec 19, 2024 · bd4af6c · bd4af6c
1 parent b0bf9aa
commit bd4af6c
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/extension/llm/modules/test/test_attention.py b/extension/llm/modules/test/test_attention.py
@@ -150,6 +150,7 @@ def test_attention_export(self):
                 (self.x, self.x),
                 kwargs={"input_pos": self.input_pos},
                 dynamic_shapes=self.dynamic_shapes,
+                strict=True,
             )
         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)
@@ -196,6 +197,7 @@ def test_attention_executorch(self):
                 (self.x, self.x),
                 kwargs={"input_pos": self.input_pos},
                 dynamic_shapes=self.dynamic_shapes,
+                strict=True,
             )
         et_program = to_edge(
             et_mha_ep,

diff --git a/extension/llm/modules/test/test_position_embeddings.py b/extension/llm/modules/test/test_position_embeddings.py
@@ -49,14 +49,14 @@ def test_tile_positional_embedding_smoke(self):
         self.assertTrue(torch.allclose(y, ref_y))
 
     def test_tile_positional_embedding_export(self):
-
         tpe_ep = torch.export.export(
             self.tpe,
             (self.x, self.aspect_ratio),
             dynamic_shapes=(
                 self.dynamic_shape,
                 None,
             ),  # assuming aspect ratio is static
+            strict=True,
         )
 
         y = tpe_ep.module()(self.x, self.aspect_ratio)
@@ -91,6 +91,7 @@ def test_tile_positional_embedding_et(self):
                 self.dynamic_shape,
                 None,
             ),  # assuming aspect ratio is static
+            strict=True,
         )
         et_program = to_edge(
             tpe_ep,
@@ -148,14 +149,14 @@ def test_tiled_token_positional_embedding_smoke(self):
         assert_close(y, ref_y)
 
     def test_tiled_token_positional_embedding_export(self):
-
         tpe_ep = torch.export.export(
             self.tpe,
             (self.x, self.aspect_ratio),
             dynamic_shapes=(
                 self.dynamic_shape,
                 None,
             ),  # assuming aspect ratio is static
+            strict=True,
         )
 
         y = tpe_ep.module()(self.x, self.aspect_ratio)
@@ -172,6 +173,7 @@ def test_tiled_token_positional_embedding_aoti(self):
                 self.dynamic_shape,
                 None,
             ),  # assuming aspect ratio is static
+            strict=True,
         )
 
         with tempfile.TemporaryDirectory() as tmpdir:
@@ -195,6 +197,7 @@ def test_tiled_token_positional_embedding_et(self):
                 self.dynamic_shape,
                 None,
             ),  # assuming aspect ratio is static
+            strict=True,
         )
         et_program = to_edge(
             tpe_ep,