executorch/backends/apple/coreml/test

Differential Revision: D70575360

Pull Request resolved: #9054

Author: gmagogsfm

backends/apple/coreml/test/test_coreml_partitioner.py

@@ -117,7 +117,7 @@ def forward(self, q, k, v, mask):
         v = torch.randn(batch_size, n_heads, max_seq_length, embedding_dim)
         mask = torch.randn(seq_len, max_seq_length)
         example_inputs = (q, k, v, mask)
-        ep = torch.export.export(model, example_inputs)
+        ep = torch.export.export(model, example_inputs, strict=True)
         coreml_partitioner = CoreMLPartitioner()
 
         # Using to_edge_transform_and_lower, we expect SDPA will be preserved and show up in delegated graph

backends/qualcomm/_passes/decompose_any.py

@@ -41,7 +41,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
                 keepdim = node.args[2] if len(node.args) > 2 else False
                 model = Any(dim, keepdim)
                 edge_mgr = to_edge(
-                    torch.export.export(model, (node.args[0].meta["val"],))
+                    torch.export.export(model, (node.args[0].meta["val"],), strict=True)
                 )
                 decomposed_module = edge_mgr.exported_program()
 

backends/qualcomm/_passes/decompose_linalg_vector_norm.py

@@ -46,11 +46,13 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
                 model = LinalgVectorNorm(ord, dim, keepdim)
                 if self.aten_dialect_capture:
                     decomposed_module = torch.export.export(
-                        model, (node.args[0].meta["val"],)
+                        model, (node.args[0].meta["val"],), strict=True
                     ).module()
                 else:
                     edge_mgr = to_edge(
-                        torch.export.export(model, (node.args[0].meta["val"],))
+                        torch.export.export(
+                            model, (node.args[0].meta["val"],), strict=True
+                        )
                     )
                     decomposed_module = edge_mgr.exported_program()
 

backends/qualcomm/tests/utils.py

@@ -526,7 +526,9 @@ def get_qdq_module(
         dynamic_shapes: Dict = None,
         bypass_check: bool = False,
     ) -> torch.fx.GraphModule:
-        m = torch.export.export(module, inputs, dynamic_shapes=dynamic_shapes).module()
+        m = torch.export.export(
+            module, inputs, dynamic_shapes=dynamic_shapes, strict=True
+        ).module()
 
         quantizer = QnnQuantizer()
         quantizer.add_custom_quant_annotations(custom_quant_annotations)

backends/qualcomm/utils/utils.py

@@ -452,7 +452,7 @@ def capture_program(
     dynamic_shapes: Dict = None,
 ) -> exir.ExirExportedProgram:
     module = _preprocess_module(module, inputs)
-    ep = torch.export.export(module, inputs, dynamic_shapes=dynamic_shapes)
+    ep = torch.export.export(module, inputs, dynamic_shapes=dynamic_shapes, strict=True)
     decomposed_ep = ep.run_decompositions(get_decomp_table())
     core_ep = ExirExportedProgram(decomposed_ep, False)
     core_ep.transform(TensorI64toI32(edge_program=core_ep))

backends/transforms/test/test_rank_0_to_rank_1.py

@@ -17,7 +17,7 @@ def forward(self, x, y):
         model.eval()
 
         example_inputs = (torch.tensor(1.0), torch.tensor(2.0))
-        aten = torch.export.export(model, example_inputs)
+        aten = torch.export.export(model, example_inputs, strict=True)
 
         # Check that the input rank is 0
         for node in aten.graph.nodes:

examples/apple/coreml/llama/export.py

@@ -192,10 +192,7 @@ def main() -> None:
     )
     example_inputs = input_manager.get_inputs(tokens=[0])
 
-    ep = torch.export.export(
-        model,
-        example_inputs,
-    )
+    ep = torch.export.export(model, example_inputs, strict=True)
     print("Exported program")
     print(ep)
 

exir/tests/test_memory_planning.py

@@ -749,7 +749,7 @@ def forward(self, input, label):
         net = TrainingNet(Net())
         inputs = (torch.randn(1, 6, 5, 5), torch.ones(1, dtype=torch.int64))
 
-        ep = export(net, inputs)
+        ep = export(net, inputs, strict=True)
         ep = _export_forward_backward(ep)
         ep = to_edge(ep)
         ep = ep.to_executorch()

exir/verification/test/test_verifier.py

@@ -153,7 +153,7 @@ def forward(self, input, label):
         net = TrainingNet(Net())
         inputs = (torch.randn(1, 6, 5, 5), torch.ones(1, dtype=torch.int64))
 
-        export_model = export(net, inputs)
+        export_model = export(net, inputs, strict=True)
         export_model = _export_forward_backward(export_model)
 
         edge = to_edge(export_model)