executorch/exir/tests

Reviewed By: avikchaudhuri

Differential Revision: D67383351

Author: gmagogsfm

exir/tests/models.py

@@ -173,10 +173,7 @@ def get_random_inputs(self) -> Tuple[Tensor, Tensor]:
 
         delegated_m = DelegateAdd()
         edge_ir_m = to_edge(
-            export(
-                delegated_m,
-                delegated_m.get_random_inputs(),
-            )
+            export(delegated_m, delegated_m.get_random_inputs(), strict=True)
         )
         lowered_module = LoweredBackendModule(
             edge_program=edge_ir_m.exported_program(),

exir/tests/test_arg_validator.py

@@ -31,7 +31,7 @@ def forward(self, x):
 
         m = TestModel()
         inputs = (torch.randn(1, 3, 100, 100).to(dtype=torch.int),)
-        egm = to_edge(export(m, inputs)).exported_program().graph_module
+        egm = to_edge(export(m, inputs, strict=True)).exported_program().graph_module
         validator = EdgeOpArgValidator(egm)
         validator.run(*inputs)
         self.assertEqual(len(validator.violating_ops), 0)
@@ -49,7 +49,7 @@ def forward(self, x):
         inputs = (torch.randn(1, 3, 100, 100).to(dtype=torch.bfloat16),)
         egm = (
             to_edge(
-                export(M(), inputs),
+                export(M(), inputs, strict=True),
                 compile_config=EdgeCompileConfig(_check_ir_validity=False),
             )
             .exported_program()

exir/tests/test_delegate.py

@@ -45,7 +45,7 @@ def g(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
             return x + y
 
         inputs = (torch.ones(1, 3), torch.ones(1, 3))
-        edge_ir_m = to_edge(export(WrapperModule(g), inputs))
+        edge_ir_m = to_edge(export(WrapperModule(g), inputs, strict=True))
         lowered_module: LoweredBackendModule = LoweredBackendModule(
             edge_ir_m.exported_program(), "BackendWithCompilerDemo", b"moo", []
         )
@@ -54,10 +54,7 @@ def f(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
             return torch.ops.higher_order.executorch_call_delegate(lowered_module, x, y)
 
         orig_res = f(*inputs)
-        gm = export(
-            WrapperModule(f),
-            inputs,
-        )
+        gm = export(WrapperModule(f), inputs, strict=True)
         FileCheck().check("lowered_module_0").check(
             "torch.ops.higher_order.executorch_call_delegate"
         ).run(gm.graph_module.code)
@@ -69,7 +66,7 @@ def test_to_backend(self) -> None:
         m = models.CompositeDelegateModule()
 
         exec_prog = to_edge(
-            export(m, m.get_random_inputs()),
+            export(m, m.get_random_inputs(), strict=True),
             compile_config=EdgeCompileConfig(_check_ir_validity=False),
         ).to_executorch()  # TODO(larryliu): fix split_copy.Tensor
         graph_module = exec_prog.exported_program().graph_module
@@ -165,7 +162,7 @@ def forward(self, x, y):
                 return x
 
         orig_res = Model()(*inputs)
-        prog = to_edge(export(Model(), inputs))
+        prog = to_edge(export(Model(), inputs, strict=True))
         gm = prog.exported_program().graph_module
 
         node_list = []
@@ -225,7 +222,7 @@ def forward(self, x, y):
                 return x
 
         orig_res = Model()(*inputs)
-        prog = to_edge(export(Model(), inputs))
+        prog = to_edge(export(Model(), inputs, strict=True))
         gm = prog.exported_program().graph_module
 
         node_list = []
@@ -284,7 +281,7 @@ def forward(self, x, y):
                 return x
 
         orig_res = Model()(*inputs)
-        prog = to_edge(export(Model(), inputs))
+        prog = to_edge(export(Model(), inputs, strict=True))
         gm = prog.exported_program().graph_module
 
         node_list = []

exir/tests/test_dynamic_shape_propagation.py

@@ -22,7 +22,12 @@ def test_repeat(self):
         inputs = inputs[0], inputs[1]
 
         prog = to_edge(
-            export(eager_model, inputs, dynamic_shapes=eager_model.get_dynamic_shape()),
+            export(
+                eager_model,
+                inputs,
+                dynamic_shapes=eager_model.get_dynamic_shape(),
+                strict=True,
+            ),
             compile_config=exir.EdgeCompileConfig(_check_ir_validity=False),
         )
 
@@ -48,7 +53,7 @@ def test_unbacked_symint(self):
         inputs = inputs[0], inputs[1]
 
         prog = to_edge(
-            export(eager_model, inputs, dynamic_shapes=None),
+            export(eager_model, inputs, dynamic_shapes=None, strict=True),
             compile_config=exir.EdgeCompileConfig(_check_ir_validity=False),
         )
         new_prog = prog.transform([SpecPropPass(), HintBasedSymShapeEvalPass()])

exir/tests/test_memory_format_ops_pass.py

@@ -269,7 +269,7 @@ def call_operator(self, op, args, kwargs, meta):
         _to_dim_order_op_str = "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default"
 
         before_epm = to_edge(
-            export(toy_model, sample_input),
+            export(toy_model, sample_input, strict=True),
             compile_config=EdgeCompileConfig(_skip_dim_order=False),
         )
 

exir/tests/test_memory_format_ops_pass_utils.py

@@ -104,7 +104,9 @@ class MemoryFormatOpsPassTestUtils:
     def memory_format_test_runner(
         test_class: unittest.TestCase, test_set: MemoryFormatTestSet
     ):
-        before = export(test_set.module, test_set.sample_input).run_decompositions({})
+        before = export(
+            test_set.module, test_set.sample_input, strict=True
+        ).run_decompositions({})
 
         if test_set.use_xnnpack:
             epm = to_edge_transform_and_lower(

exir/tests/test_memory_planning.py

@@ -239,12 +239,7 @@ def wrapper(self: "TestMemoryPlanning") -> None:
             #  torch._tensor.Tensor]` is not a function.
             inputs = eager_module.get_random_inputs()
             graph_module = (
-                to_edge(
-                    export(
-                        eager_module,
-                        inputs,
-                    )
-                )
+                to_edge(export(eager_module, inputs, strict=True))
                 .exported_program()
                 .graph_module
             )
@@ -491,10 +486,7 @@ def test_multiple_pools(
         expected_bufsizes: List[int],
     ) -> None:
         edge_program = to_edge(
-            export(
-                MultiplePoolsToyModel(),
-                (torch.ones(1),),
-            )
+            export(MultiplePoolsToyModel(), (torch.ones(1),), strict=True)
         )
 
         edge_program.to_executorch(
@@ -538,7 +530,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
                 return torch.nn.functional.sigmoid(self.linear(x) + self.constant + 1)
 
         def count_planned_inputs(
-            nodes: List[Node], graph_signature: Any  # pyre-ignore
+            nodes: List[Node],
+            graph_signature: Any,  # pyre-ignore
         ) -> Tuple[int, int]:
             num_mem_planned_placeholders = 0
             num_placeholders = 0
@@ -555,7 +548,9 @@ def count_planned_inputs(
         model = Simple()
         inputs = (torch.randn(5, 5),)
 
-        ep_no_input_planning = to_edge(export(model, inputs)).to_executorch(
+        ep_no_input_planning = to_edge(
+            export(model, inputs, strict=True)
+        ).to_executorch(
             config=ExecutorchBackendConfig(
                 memory_planning_pass=MemoryPlanningPass(alloc_graph_input=False),
                 sym_shape_eval_pass=ConstraintBasedSymShapeEvalPass(),
@@ -575,7 +570,7 @@ def count_planned_inputs(
             5,  # x, self.constant, linear weight, linear bias, '1' scalar promoted to tensor
         )
 
-        ep_input_planning = to_edge(export(model, inputs)).to_executorch(
+        ep_input_planning = to_edge(export(model, inputs, strict=True)).to_executorch(
             config=ExecutorchBackendConfig(
                 memory_planning_pass=MemoryPlanningPass(alloc_graph_input=True),
                 sym_shape_eval_pass=ConstraintBasedSymShapeEvalPass(),
@@ -609,7 +604,7 @@ def forward(self, a, b, x):
 
         model = TestModel()
         example_inputs = (torch.rand(1, 6, 2), torch.rand(1, 6, 2), torch.randn(5, 5))
-        exported_model = torch.export.export(model, example_inputs)
+        exported_model = torch.export.export(model, example_inputs, strict=True)
         edge = to_edge(exported_model)
 
         class TestPass(ExportPass):