executorch/exir/program/test

Differential Revision: D67383235

Pull Request resolved: #7397

Author: gmagogsfm

backends/apple/coreml/runtime/test/export_stateful_model.py

@@ -47,7 +47,7 @@ def main() -> None:
         torch.randn((1, embedding_dim)),
         torch.tensor([0]),
     )
-    exported_model = export(model, example_inputs)
+    exported_model = export(model, example_inputs, strict=True)
     edge_program_manager = exir.to_edge(exported_model)
     compile_specs = CoreMLBackend.generate_compile_specs(
         compute_precision=ct.precision.FLOAT16,

backends/apple/coreml/test/test_coreml_partitioner.py

@@ -16,7 +16,6 @@
 
 
 class TestCoreMLPartitioner(unittest.TestCase):
-
     # TODO(T182928844): Delegate dim order op to backend.
     edge_compile_config = executorch.exir.EdgeCompileConfig(_skip_dim_order=True)
 
@@ -34,7 +33,7 @@ def forward(self, a, x, b):
         model.eval()
 
         example_inputs = (torch.randn(2, 2), torch.randn(2, 2), torch.randn(2, 2))
-        exir_program_aten = torch.export.export(model, example_inputs)
+        exir_program_aten = torch.export.export(model, example_inputs, strict=True)
 
         edge_program_manager = executorch.exir.to_edge(
             exir_program_aten, compile_config=self.edge_compile_config
@@ -61,7 +60,7 @@ def test_vit_skip_conv(self):
         model.eval()
 
         example_inputs = (torch.randn(1, 3, 224, 224),)
-        exir_program_aten = torch.export.export(model, example_inputs)
+        exir_program_aten = torch.export.export(model, example_inputs, strict=True)
         edge_program_manager = executorch.exir.to_edge(
             exir_program_aten, compile_config=self.edge_compile_config
         )
@@ -106,7 +105,7 @@ def forward(self, q, k_val, input_pos):
         k_val = torch.randn((1, embedding_dim))
         input_pos = torch.tensor([0])
         example_inputs = (q, k_val, input_pos)
-        exir_program_aten = torch.export.export(model, example_inputs)
+        exir_program_aten = torch.export.export(model, example_inputs, strict=True)
 
         compile_specs = CoreMLBackend.generate_compile_specs(
             minimum_deployment_target=ct.target.iOS18

backends/apple/mps/test/test_mps_utils.py

@@ -247,10 +247,7 @@ def lower_module_and_test_output(
             )
 
             executorch_program = to_edge(
-                export(
-                    delegated_program,
-                    sample_inputs,
-                ),
+                export(delegated_program, sample_inputs, strict=True),
                 compile_config=exir.EdgeCompileConfig(
                     _check_ir_validity=False,
                     _skip_dim_order=True,  # TODO(T182928844): Delegate dim order op to backend.

backends/cadence/aot/compiler.py

@@ -176,7 +176,7 @@ def export_program(
     torch._C._set_mkldnn_enabled(False)
 
     # else: capture the model and return it.
-    expo_program = export(model, inputs)
+    expo_program = export(model, inputs, strict=True)
 
     if dump_graphs:
         logging.info("Exported graph:")

backends/example/test_example_delegate.py

@@ -60,7 +60,7 @@ def get_example_inputs():
 
         quantized_gm = m
         exported_program = to_edge(
-            export(quantized_gm, copy.deepcopy(example_inputs)),
+            export(quantized_gm, copy.deepcopy(example_inputs), strict=True),
             compile_config=EDGE_COMPILE_CONFIG,
         )
 
@@ -92,7 +92,7 @@ def test_delegate_mobilenet_v2(self):
 
         quantized_gm = m
         exported_program = to_edge(
-            export(quantized_gm, copy.deepcopy(example_inputs)),
+            export(quantized_gm, copy.deepcopy(example_inputs), strict=True),
             compile_config=EDGE_COMPILE_CONFIG,
         )
 

backends/qualcomm/tests/test_qnn_delegate.py

@@ -1617,7 +1617,7 @@ def test_qnn_backend_multi_contexts_composite(self):
         )
         sample_input = module.get_random_input()
         edge_prog = to_edge(
-            torch.export.export(module, sample_input),
+            torch.export.export(module, sample_input, strict=True),
         )
         update_spill_fill_size(edge_prog.exported_program())
         exec_prog = edge_prog.to_executorch()
@@ -1957,7 +1957,7 @@ def calibrator(gm):
         self.assertEqual(len(exported_progs), 1)
         # lower all graph again, the skipped operators will be left in CPU
         exec_prog = to_edge(
-            torch.export.export(graph_module, sample_input),
+            torch.export.export(graph_module, sample_input, strict=True),
         ).to_executorch()
         self.verify_output(module, sample_input, exec_prog)
 
@@ -2004,7 +2004,7 @@ def calibrator(gm):
         self.assertEqual(len(exported_progs), 2)
         # lower all graph again, the skipped operators will be left in CPU
         exec_prog = exec_prog = to_edge(
-            torch.export.export(graph_module, sample_input),
+            torch.export.export(graph_module, sample_input, strict=True),
         ).to_executorch()
         self.verify_output(module, sample_input, exec_prog)
 
@@ -2041,7 +2041,7 @@ def calibrator(gm):
         self.assertEqual(len(exported_progs), 5)
         # lower all graph again, the skipped operators will be delegated with fp16
         exec_prog = to_edge(
-            torch.export.export(graph_module, sample_input),
+            torch.export.export(graph_module, sample_input, strict=True),
         ).to_executorch()
         self.verify_output(module, sample_input, exec_prog)
 
@@ -2086,7 +2086,7 @@ def test_qnn_backend_multi_contexts_composite(self):
         )
         sample_input = module.get_random_input()
         edge_prog = to_edge(
-            torch.export.export(module, sample_input),
+            torch.export.export(module, sample_input, strict=True),
         )
         update_spill_fill_size(edge_prog.exported_program())
         exec_prog = edge_prog.to_executorch()
@@ -2721,7 +2721,6 @@ def test_ssd300_vgg16(self):
 
 
 class TestExampleQaihubScript(TestQNN):
-
     def required_envs(self, conditions=None) -> bool:
         conditions = [] if conditions is None else conditions
         return all(

backends/qualcomm/tests/utils.py

@@ -385,7 +385,7 @@ def get_qdq_module(
         custom_quant_annotations: Tuple[Callable] = (),
         quant_dtype: QuantDtype = QuantDtype.use_8a8w,
     ) -> torch.fx.GraphModule:
-        m = torch.export.export(module, inputs).module()
+        m = torch.export.export(module, inputs, strict=True).module()
 
         quantizer = QnnQuantizer()
         quantizer.add_custom_quant_annotations(custom_quant_annotations)

backends/qualcomm/utils/utils.py

@@ -337,7 +337,7 @@ def capture_program(
     inputs: Tuple[torch.Tensor],
     custom_pass_config: FrozenSet[str] = frozenset(),
 ) -> exir.ExirExportedProgram:
-    ep = torch.export.export(module, inputs)
+    ep = torch.export.export(module, inputs, strict=True)
     decomposed_ep = ep.run_decompositions(get_decomp_table())
     # We choose call_operator by target in ConvertBinaryOpsWithScalar
     # because it is the same source_fn_stack for MultiheadAttention
@@ -551,7 +551,7 @@ def prepare_subgm(subgm, subgm_name):
 
     fp_node_id_set = fp_node_id_set if fp_node_id_set is not None else set()
     fp_node_op_set = fp_node_op_set if fp_node_op_set is not None else set()
-    graph_module = torch.export.export(nn_module, sample_input).module()
+    graph_module = torch.export.export(nn_module, sample_input, strict=True).module()
     # define node support type
     capability_partitioner = CapabilityBasedPartitioner(
         graph_module,
@@ -664,7 +664,7 @@ def forward(self, *inputs):
                 ).default(inputs)
 
         model = Model()
-        prog = torch.export.export(model, tuple(inputs.values()))
+        prog = torch.export.export(model, tuple(inputs.values()), strict=True)
         # bookkeeping for variables' life cycle
         return {
             "custom_op": custom_op,

backends/vulkan/test/test_vulkan_delegate.py

@@ -112,7 +112,7 @@ def run_test():
             model(*sample_inputs)
 
             program: ExportedProgram = export(
-                model, sample_inputs, dynamic_shapes=dynamic_shapes
+                model, sample_inputs, dynamic_shapes=dynamic_shapes, strict=True
             )
 
             edge_program = to_edge_transform_and_lower(

backends/xnnpack/partition/graphs/sdpa.py

@@ -76,6 +76,7 @@ def forward(
                         v,
                         mask,
                     ),
+                    strict=True,
                 ),
                 compile_config=get_xnnpack_edge_compile_config(),
             )