Arm backend: Replace .export_for_training with .export (pytorch#13280)

Signed-off-by: Adrian Lundell <[email protected]>

Author: AdrianLundell

backends/arm/test/misc/test_extract_io_params_tosa.py

@@ -60,16 +60,12 @@ def test_roundtrip_extracts_io_params(builder_method, quantizer_cls, partitioner
     operator_config = get_symmetric_quantization_config(is_qat=True)
     quantizer.set_global(operator_config)
 
-    exported = torch.export.export_for_training(
-        mod, copy.deepcopy(example_inputs), strict=True
-    )
+    exported = torch.export.export(mod, copy.deepcopy(example_inputs), strict=True)
     prepared = prepare_pt2e(exported.module(), quantizer)
     _ = prepared(*example_inputs)
 
     converted = convert_pt2e(prepared)
-    final_export = torch.export.export_for_training(
-        converted, example_inputs, strict=True
-    )
+    final_export = torch.export.export(converted, example_inputs, strict=True)
     partitioner = partitioner_cls(compile_spec)
     edge_prog = to_edge_transform_and_lower(final_export, partitioner=[partitioner])
 

backends/cortex_m/test/test_replace_quant_nodes.py

@@ -1,5 +1,6 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
+# Copyright 2025 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
@@ -16,7 +17,7 @@
     ReplaceQuantNodesPass,
 )
 from executorch.exir.dialects._ops import ops as exir_ops
-from torch.export import export, export_for_training
+from torch.export import export
 from torch.fx import GraphModule
 from torchao.quantization.pt2e.observer import HistogramObserver
 from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e
@@ -125,9 +126,7 @@ def forward(self, x):
         example_inputs = (torch.randn(10, 11, 12),)
 
         # Step 1: Export and quantize the model
-        exported_model = export_for_training(
-            model.eval(), example_inputs, strict=True
-        ).module()
+        exported_model = export(model.eval(), example_inputs, strict=True).module()
         prepared_model = prepare_pt2e(exported_model, AddQuantizer())
         quantized_model = convert_pt2e(prepared_model)
 

docs/source/backends-arm-ethos-u.md

@@ -50,14 +50,14 @@ compile_spec = ArmCompileSpecBuilder().ethosu_compile_spec(
     ).build()
 
 # Post training quantization
-graph_module = torch.export.export_for_training(mobilenet_v2, example_inputs).module()
+graph_module = torch.export.export(mobilenet_v2, example_inputs).module()
 quantizer = EthosUQuantizer(compile_spec)
 operator_config = get_symmetric_quantization_config(is_per_channel=False)
 quantizer.set_global(operator_config)
 graph_module = prepare_pt2e(graph_module, quantizer)
 graph_module(*example_inputs)
 graph_module = convert_pt2e(graph_module)
-exported_program = torch.export.export_for_training(graph_module, example_inputs)
+exported_program = torch.export.export(graph_module, example_inputs)
 
 # Lower the exported program to the Ethos-U backend and save pte file.
 edge_program_manager = to_edge_transform_and_lower(

examples/arm/aot_arm_compiler.py

@@ -710,7 +710,7 @@ def quantize_model(args, model: torch.nn.Module, example_inputs, compile_spec):
         args.evaluate_config,
     )
     # Wrap quantized model back into an exported_program
-    exported_program = torch.export.export_for_training(
+    exported_program = torch.export.export(
         model_int8, example_inputs, strict=args.strict_export
     )
 
@@ -803,9 +803,9 @@ def transform_for_cortex_m_backend(edge):
     )
     model = original_model.eval()
 
-    # export_for_training under the assumption we quantize, the exported form also works
+    # export under the assumption we quantize, the exported form also works
     # in to_edge if we don't quantize
-    exported_program = torch.export.export_for_training(
+    exported_program = torch.export.export(
         model, example_inputs, strict=args.strict_export
     )
     model = exported_program.module()

examples/arm/ethos_u_minimal_example.ipynb

@@ -58,7 +58,7 @@
     "\n",
     "model = Add()\n",
     "model = model.eval()\n",
-    "exported_program = torch.export.export_for_training(model, example_inputs)\n",
+    "exported_program = torch.export.export(model, example_inputs)\n",
     "graph_module = exported_program.module()\n",
     "\n",
     "_ = graph_module.print_readable()"
@@ -114,7 +114,7 @@
     "_ = quantized_graph_module.print_readable()\n",
     "\n",
     "# Create a new exported program using the quantized_graph_module\n",
-    "quantized_exported_program = torch.export.export_for_training(quantized_graph_module, example_inputs)"
+    "quantized_exported_program = torch.export.export(quantized_graph_module, example_inputs)"
    ]
   },
   {