executorch/exir/program/test (pytorch#7397)

Summary: Pull Request resolved: pytorch#7397

Reviewed By: avikchaudhuri

Differential Revision: D67383235

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,

devtools/backend_debug/tests/test_delegation_info.py

@@ -31,7 +31,7 @@ def forward(self, a, x, b):
 
         m = Model()
         inputs = (torch.randn(2, 2), torch.randn(2, 2), torch.randn(2, 2))
-        edge = to_edge(torch.export.export(m, inputs)).to_backend(
+        edge = to_edge(torch.export.export(m, inputs, strict=True)).to_backend(
             AddMulPartitionerDemo()
         )
         delegation_info = get_delegation_info(edge.exported_program().graph_module)

devtools/bundled_program/util/test_util.py

@@ -271,6 +271,7 @@ def get_common_executorch_program() -> (
         m_name: export(
             StatefulWrapperModule(eager_model, getattr(eager_model, m_name)),
             capture_inputs[m_name],
+            strict=True,
         )
         for m_name in eager_model.method_names
     }

devtools/etrecord/tests/etrecord_test.py

@@ -69,7 +69,7 @@ def get_test_model_with_bundled_program(self):
 
     def get_test_model_with_manager(self):
         f = models.BasicSinMax()
-        aten_dialect = export(f, f.get_random_inputs())
+        aten_dialect = export(f, f.get_random_inputs(), strict=True)
         edge_program: EdgeProgramManager = to_edge(
             aten_dialect, compile_config=EdgeCompileConfig(_check_ir_validity=False)
         )

docs/source/tutorials_source/devtools-integration-tutorial.py

@@ -89,10 +89,7 @@ def forward(self, x):
 
 model = Net()
 
-aten_model: ExportedProgram = export(
-    model,
-    (torch.randn(1, 1, 32, 32),),
-)
+aten_model: ExportedProgram = export(model, (torch.randn(1, 1, 32, 32),), strict=True)
 
 edge_program_manager: EdgeProgramManager = to_edge(
     aten_model, compile_config=EdgeCompileConfig(_check_ir_validity=True)
@@ -141,7 +138,7 @@ def forward(self, x):
 
 # Step 1: ExecuTorch Program Export
 m_name = "forward"
-method_graphs = {m_name: export(model, (torch.randn(1, 1, 32, 32),))}
+method_graphs = {m_name: export(model, (torch.randn(1, 1, 32, 32),), strict=True)}
 
 # Step 2: Construct Method Test Suites
 inputs = [[torch.randn(1, 1, 32, 32)] for _ in range(2)]

docs/source/tutorials_source/export-to-executorch-tutorial.py

@@ -66,7 +66,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 
 example_args = (torch.randn(1, 3, 256, 256),)
-aten_dialect: ExportedProgram = export(SimpleConv(), example_args)
+aten_dialect: ExportedProgram = export(SimpleConv(), example_args, strict=True)
 print(aten_dialect)
 
 ######################################################################
@@ -101,7 +101,7 @@ def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
 
 
 example_args = (torch.randn(3, 3), torch.randn(3, 3))
-aten_dialect: ExportedProgram = export(Basic(), example_args)
+aten_dialect: ExportedProgram = export(Basic(), example_args, strict=True)
 
 # Works correctly
 print(aten_dialect.module()(torch.ones(3, 3), torch.ones(3, 3)))
@@ -131,7 +131,7 @@ def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
 dim1_x = Dim("dim1_x", min=1, max=10)
 dynamic_shapes = {"x": {1: dim1_x}, "y": {1: dim1_x}}
 aten_dialect: ExportedProgram = export(
-    Basic(), example_args, dynamic_shapes=dynamic_shapes
+    Basic(), example_args, dynamic_shapes=dynamic_shapes, strict=True
 )
 print(aten_dialect)
 
@@ -213,7 +213,7 @@ def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
 print("Quantized Graph")
 print(converted_graph)
 
-aten_dialect: ExportedProgram = export(converted_graph, example_args)
+aten_dialect: ExportedProgram = export(converted_graph, example_args, strict=True)
 print("ATen Dialect Graph")
 print(aten_dialect)
 
@@ -243,7 +243,7 @@ def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
 from executorch.exir import EdgeProgramManager, to_edge
 
 example_args = (torch.randn(1, 3, 256, 256),)
-aten_dialect: ExportedProgram = export(SimpleConv(), example_args)
+aten_dialect: ExportedProgram = export(SimpleConv(), example_args, strict=True)
 
 edge_program: EdgeProgramManager = to_edge(aten_dialect)
 print("Edge Dialect Graph")
@@ -267,10 +267,10 @@ def forward(self, x):
 
 
 encode_args = (torch.randn(1, 10),)
-aten_encode: ExportedProgram = export(Encode(), encode_args)
+aten_encode: ExportedProgram = export(Encode(), encode_args, strict=True)
 
 decode_args = (torch.randn(1, 5),)
-aten_decode: ExportedProgram = export(Decode(), decode_args)
+aten_decode: ExportedProgram = export(Decode(), decode_args, strict=True)
 
 edge_program: EdgeProgramManager = to_edge(
     {"encode": aten_encode, "decode": aten_decode}
@@ -291,7 +291,7 @@ def forward(self, x):
 # rather than the ``torch.ops.aten`` namespace.
 
 example_args = (torch.randn(1, 3, 256, 256),)
-aten_dialect: ExportedProgram = export(SimpleConv(), example_args)
+aten_dialect: ExportedProgram = export(SimpleConv(), example_args, strict=True)
 edge_program: EdgeProgramManager = to_edge(aten_dialect)
 print("Edge Dialect Graph")
 print(edge_program.exported_program())
@@ -357,7 +357,7 @@ def forward(self, x):
 
 # Export and lower the module to Edge Dialect
 example_args = (torch.ones(1),)
-aten_dialect: ExportedProgram = export(LowerableModule(), example_args)
+aten_dialect: ExportedProgram = export(LowerableModule(), example_args, strict=True)
 edge_program: EdgeProgramManager = to_edge(aten_dialect)
 to_be_lowered_module = edge_program.exported_program()
 
@@ -423,7 +423,7 @@ def forward(self, x):
 
 
 example_args = (torch.ones(1),)
-aten_dialect: ExportedProgram = export(ComposedModule(), example_args)
+aten_dialect: ExportedProgram = export(ComposedModule(), example_args, strict=True)
 edge_program: EdgeProgramManager = to_edge(aten_dialect)
 exported_program = edge_program.exported_program()
 print("Edge Dialect graph")
@@ -461,7 +461,7 @@ def forward(self, a, x, b):
 
 
 example_args = (torch.randn(2, 2), torch.randn(2, 2), torch.randn(2, 2))
-aten_dialect: ExportedProgram = export(Foo(), example_args)
+aten_dialect: ExportedProgram = export(Foo(), example_args, strict=True)
 edge_program: EdgeProgramManager = to_edge(aten_dialect)
 exported_program = edge_program.exported_program()
 print("Edge Dialect graph")
@@ -495,7 +495,7 @@ def forward(self, a, x, b):
 
 
 example_args = (torch.randn(2, 2), torch.randn(2, 2), torch.randn(2, 2))
-aten_dialect: ExportedProgram = export(Foo(), example_args)
+aten_dialect: ExportedProgram = export(Foo(), example_args, strict=True)
 edge_program: EdgeProgramManager = to_edge(aten_dialect)
 exported_program = edge_program.exported_program()
 delegated_program = edge_program.to_backend(AddMulPartitionerDemo())
@@ -577,7 +577,9 @@ def forward(self, x):
 pre_autograd_aten_dialect = export_for_training(M(), example_args).module()
 # Optionally do quantization:
 # pre_autograd_aten_dialect = convert_pt2e(prepare_pt2e(pre_autograd_aten_dialect, CustomBackendQuantizer))
-aten_dialect: ExportedProgram = export(pre_autograd_aten_dialect, example_args)
+aten_dialect: ExportedProgram = export(
+    pre_autograd_aten_dialect, example_args, strict=True
+)
 edge_program: exir.EdgeProgramManager = exir.to_edge(aten_dialect)
 # Optionally do delegation:
 # edge_program = edge_program.to_backend(CustomBackendPartitioner)

examples/apple/coreml/scripts/export.py

@@ -88,7 +88,9 @@ def partition_module_to_coreml(module):
 
 def lower_module_to_coreml(module, compile_specs, example_inputs):
     module = module.eval()
-    edge = to_edge(export(module, example_inputs), compile_config=_EDGE_COMPILE_CONFIG)
+    edge = to_edge(
+        export(module, example_inputs, strict=True), compile_config=_EDGE_COMPILE_CONFIG
+    )
     # All of the subsequent calls on the edge_dialect_graph generated above (such as delegation or
     # to_executorch()) are done in place and the graph is also modified in place. For debugging purposes
     # we would like to keep a copy of the original edge dialect graph and hence we create a deepcopy of
@@ -107,7 +109,8 @@ def lower_module_to_coreml(module, compile_specs, example_inputs):
 def export_lowered_module_to_executorch_program(lowered_module, example_inputs):
     lowered_module(*example_inputs)
     exec_prog = to_edge(
-        export(lowered_module, example_inputs), compile_config=_EDGE_COMPILE_CONFIG
+        export(lowered_module, example_inputs, strict=True),
+        compile_config=_EDGE_COMPILE_CONFIG,
     ).to_executorch(config=exir.ExecutorchBackendConfig(extract_delegate_segments=True))
 
     return exec_prog
@@ -170,7 +173,7 @@ def main():
 
     if args.use_partitioner:
         model.eval()
-        exir_program_aten = torch.export.export(model, example_inputs)
+        exir_program_aten = torch.export.export(model, example_inputs, strict=True)
 
         edge_program_manager = exir.to_edge(exir_program_aten)
         edge_copy = copy.deepcopy(edge_program_manager)

examples/apple/coreml/scripts/inspector_utils.py

@@ -65,9 +65,7 @@ def build_devtools_runner_including_coreml(
     build_devtools_runner_command: str = (
         "./examples/devtools/build_example_runner.sh --coreml"
     )
-    build_command: str = (
-        f"{cd_root_command} && {conda_activate_env_command} && {build_devtools_runner_command}"
-    )
+    build_command: str = f"{cd_root_command} && {conda_activate_env_command} && {build_devtools_runner_command}"
     subprocess.run(
         f'bash -c "{build_command}"', shell=True, check=True
     ).check_returncode()
@@ -87,10 +85,7 @@ def to_core_aten(
     module: torch.nn.Module,
     example_inputs: Tuple[Value, ...],
 ) -> ExportedProgram:
-    core_aten_program = export(
-        mod=module,
-        args=example_inputs,
-    )
+    core_aten_program = export(mod=module, args=example_inputs, strict=True)
     return core_aten_program
 
 

examples/devtools/scripts/gen_sample_etrecord.py

@@ -31,10 +31,7 @@
 
 def gen_etrecord(model: torch.nn.Module, inputs: Any, output_path=None):
     f = model
-    aten_dialect: ExportedProgram = export(
-        f,
-        inputs,
-    )
+    aten_dialect: ExportedProgram = export(f, inputs, strict=True)
     edge_program: EdgeProgramManager = to_edge(
         aten_dialect, compile_config=EdgeCompileConfig(_check_ir_validity=True)
     )

examples/llm_manual/export_nanogpt.py

@@ -30,7 +30,7 @@
     m = export_for_training(
         model, example_inputs, dynamic_shapes=dynamic_shape
     ).module()
-    traced_model = export(m, example_inputs, dynamic_shapes=dynamic_shape)
+    traced_model = export(m, example_inputs, dynamic_shapes=dynamic_shape, strict=True)
 
 # Convert the model into a runnable ExecuTorch program.
 # To be further lowered to Xnnpack backend, `traced_model` needs xnnpack-specific edge compile config

examples/mediatek/aot_utils/oss_utils/utils.py

@@ -37,9 +37,9 @@ def build_executorch_binary(
         for data in dataset:
             annotated_model(*data)
         quantized_model = convert_pt2e(annotated_model, fold_quantize=False)
-        aten_dialect = torch.export.export(quantized_model, inputs)
+        aten_dialect = torch.export.export(quantized_model, inputs, strict=True)
     else:
-        aten_dialect = torch.export.export(model, inputs)
+        aten_dialect = torch.export.export(model, inputs, strict=True)
 
     from executorch.exir.program._program import to_edge_transform_and_lower
 

examples/mediatek/model_export_scripts/llama.py

@@ -338,7 +338,7 @@ def export_to_et_ir(
         print(f"Exporting Shape {shape} to:\n{dest_path}")
         example_inputs = model.get_example_inputs(*ntok_and_cache)
         aten_dialect: exir.ExportedProgram = torch.export.export(
-            converted_graph, example_inputs
+            converted_graph, example_inputs, strict=True
         )
 
         print("Lowering to Edge Dialect Graph")

examples/models/llama3_2_vision/text_decoder/test/test_text_decoder.py

@@ -70,6 +70,7 @@ def test_llama3_2_text_decoder_aoti(self) -> None:
                 model.get_example_inputs(),
                 kwargs=model.get_example_kwarg_inputs(),
                 dynamic_shapes=model.get_dynamic_shapes(),
+                strict=True,
             )
         with tempfile.TemporaryDirectory() as tmpdir:
             path = torch._inductor.aoti_compile_and_package(

examples/models/llama3_2_vision/vision_encoder/test/test_vision_encoder.py

@@ -32,6 +32,7 @@ def test_flamingo_vision_encoder(self) -> None:
             encoder,
             model.get_example_inputs(),
             dynamic_shapes=model.get_dynamic_shapes(),
+            strict=True,
         )
         with tempfile.TemporaryDirectory() as tmpdir:
             path = torch._inductor.aoti_compile_and_package(

examples/models/llava/export_llava.py

@@ -116,6 +116,7 @@ def forward(self, input_pos, embeddings):
             manager.pre_autograd_graph_module,
             manager.example_inputs,
             dynamic_shapes=manager._get_dynamic_shape(),
+            strict=True,
         )
     return text_model_ep
 
@@ -158,6 +159,7 @@ def forward(self, images):
             manager.pre_autograd_graph_module,
             manager.example_inputs,
             dynamic_shapes=manager.dynamic_shapes,
+            strict=True,
         )
     return image_encoder_ep
 
@@ -176,7 +178,10 @@ def quant_embedding(model):
     dynamic_shapes = [{1: token_dim_1}]
     with torch.no_grad():
         token_embedding_ep = torch.export.export(
-            quantized_token_embed.embed_tokens, (prompt,), dynamic_shapes=dynamic_shapes
+            quantized_token_embed.embed_tokens,
+            (prompt,),
+            dynamic_shapes=dynamic_shapes,
+            strict=True,
         )
     return token_embedding_ep
 

examples/models/phi-3-mini-lora/export_model.py

@@ -55,7 +55,7 @@ def export_phi3_mini_lora(model) -> None:
     tokens = randint(0, vocab_size, (batch_size, seq_len), dtype=long)
     example_args = (tokens,)
     with sdpa_kernel([SDPBackend.MATH]):
-        aten_dialect: ExportedProgram = export(model, example_args)
+        aten_dialect: ExportedProgram = export(model, example_args, strict=True)
 
         # 2. to_edge: Make optimizations for Edge devices.
         print("Lowering to edge dialect")
@@ -93,7 +93,7 @@ def export_phi3_mini_lora_training(model) -> None:
     labels = tokens
     example_args = (tokens, labels)
     with sdpa_kernel([SDPBackend.MATH]):
-        exported_graph: ExportedProgram = export(model, example_args)
+        exported_graph: ExportedProgram = export(model, example_args, strict=True)
         print("Creating a joint forward-backwards graph for training")
         joint_graph = _export_forward_backward(exported_graph)
 

examples/qualcomm/oss_scripts/llama2/llama.py

@@ -108,7 +108,6 @@ def annotate_cat(node: Node, quantization_config: QuantizationConfig):
     def annotate_single_in_single_out(
         node: Node, quantization_config: QuantizationConfig
     ) -> None:
-
         input_qspec_map = {}
         input_act = node.args[0]
         input_qspec_map[input_act] = quantization_config.input_activation
@@ -356,7 +355,7 @@ def quantize(self, quant_dtype, custom_annotations=()):
 
         with torch.no_grad():
             fx_graph_module = torch.export.export(
-                self.llama_model, self.inputs
+                self.llama_model, self.inputs, strict=True
             ).module()
             fx_graph_module = prepare_pt2e(fx_graph_module, quantizer)
         print("Quantizing the model...")

examples/qualcomm/oss_scripts/llama3_2/llama.py

@@ -236,7 +236,7 @@ def quantize(self, quant_dtype, args, custom_annotations=()):
 
         with torch.no_grad():
             fx_graph_module = torch.export.export(
-                self.llama_model, self.inputs
+                self.llama_model, self.inputs, strict=True
             ).module()
             fx_graph_module = prepare_pt2e(fx_graph_module, quantizer)
         logging.info("Quantizing the model...")

examples/qualcomm/scripts/export_example.py

@@ -61,7 +61,7 @@ def main() -> None:
     quantizer = QnnQuantizer()
 
     # Typical pytorch 2.0 quantization flow
-    m = torch.export.export(model.eval(), example_inputs).module()
+    m = torch.export.export(model.eval(), example_inputs, strict=True).module()
     m = prepare_pt2e(m, quantizer)
     # Calibration
     m(*example_inputs)

examples/qualcomm/scripts/mobilebert_fine_tune.py

@@ -292,7 +292,7 @@ def calibrator(gm):
         )
         # lower all graph again, the skipped operators will be left in CPU
         exec_prog = to_edge(
-            torch.export.export(graph_module, inputs[0]),
+            torch.export.export(graph_module, inputs[0], strict=True),
         ).to_executorch()
 
         with open(f"{args.artifact}/{pte_filename}.pte", "wb") as file:

examples/qualcomm/utils.py

@@ -281,7 +281,7 @@ def build_executorch_binary(
         None: The function writes the output to a specified .pte file.
     """
     if quant_dtype is not None:
-        captured_model = torch.export.export(model, inputs).module()
+        captured_model = torch.export.export(model, inputs, strict=True).module()
         if qat_training_data:
             quantizer = custom_quantizer or make_quantizer(
                 quant_dtype=quant_dtype, is_qat=True