NXP backend: Update deprecated export call.

Author: MartinPavella

backends/nxp/tests/executorch_pipeline.py

@@ -4,8 +4,6 @@
 # LICENSE file in the root directory of this source tree.
 
 import torch
-from torch import nn
-from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e
 
 from executorch import exir
 from executorch.backends.nxp.backend.custom_delegation_options import (
@@ -27,6 +25,8 @@
     ExecutorchProgramManager,
 )
 from executorch.extension.export_util.utils import export_to_edge
+from torch import nn
+from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e
 
 
 def _quantize_model(model, calibration_inputs: list[tuple[torch.Tensor]]):
@@ -73,9 +73,7 @@ def to_quantized_edge_program(
     # Make sure the model is in the evaluation mode.
     model.eval()
 
-    exir_program_aten = torch.export.export_for_training(
-        model, example_input, strict=True
-    )
+    exir_program_aten = torch.export.export(model, example_input, strict=True)
 
     exir_program_aten__module_quant = _quantize_model(
         exir_program_aten.module(), calibration_inputs

backends/nxp/tests/test_batch_norm_fusion.py

@@ -95,7 +95,7 @@ def test_batch_norm_conv_fusing(bias: bool, input_shape: list[int]):
     example_input = (torch.ones(*input_shape),)
 
     module = ConvBatchNormModule(bias, len(input_shape), 4)
-    program = torch.export.export_for_training(module, example_input, strict=True)
+    program = torch.export.export(module, example_input, strict=True)
     og_module = program.module()
 
     pm = NeutronAtenPassManager()
@@ -129,7 +129,7 @@ def test_batch_norm_linear_fusing(bias: bool):
     example_input = (torch.ones(*input_shape),)
 
     module = LinearBatchNormModule(bias, 4, input_shape[-1], input_shape[1])
-    program = torch.export.export_for_training(module, example_input, strict=True)
+    program = torch.export.export(module, example_input, strict=True)
     og_module = program.module()
 
     pm = NeutronAtenPassManager()

backends/nxp/tests/test_quantizer.py

@@ -23,9 +23,7 @@ def test_quantizer_conv2d():
 
     example_input = (torch.ones(1, 4, 32, 32),)
     quantizer = NeutronQuantizer()
-    graph_module = torch.export.export_for_training(
-        model, example_input, strict=True
-    ).module()
+    graph_module = torch.export.export(model, example_input, strict=True).module()
 
     # noinspection PyTypeChecker
     m = prepare_pt2e(graph_module, quantizer)
@@ -64,9 +62,7 @@ def test_quantizer_linear():
 
     example_input = (torch.ones(10, 32),)
     quantizer = NeutronQuantizer()
-    graph_module = torch.export.export_for_training(
-        model, example_input, strict=True
-    ).module()
+    graph_module = torch.export.export(model, example_input, strict=True).module()
 
     # noinspection PyTypeChecker
     m = prepare_pt2e(graph_module, quantizer)
@@ -105,9 +101,7 @@ def test_quantizer_maxpool2d():
 
     example_input = (torch.ones(1, 8, 32, 32),)
     quantizer = NeutronQuantizer()
-    graph_module = torch.export.export_for_training(
-        model, example_input, strict=True
-    ).module()
+    graph_module = torch.export.export(model, example_input, strict=True).module()
 
     # noinspection PyTypeChecker
     m = prepare_pt2e(graph_module, quantizer)
@@ -143,9 +137,7 @@ def test_quantizer_softmax():
 
     example_input = (torch.ones(1, 10),)
     quantizer = NeutronQuantizer()
-    graph_module = torch.export.export_for_training(
-        model, example_input, strict=True
-    ).module()
+    graph_module = torch.export.export(model, example_input, strict=True).module()
 
     # noinspection PyTypeChecker
     m = prepare_pt2e(graph_module, quantizer)
@@ -182,9 +174,7 @@ def test_quantizer_single_maxpool2d():
 
     example_input = (torch.ones(1, 4, 32, 32),)
     quantizer = NeutronQuantizer()
-    graph_module = torch.export.export_for_training(
-        model, example_input, strict=True
-    ).module()
+    graph_module = torch.export.export(model, example_input, strict=True).module()
 
     # noinspection PyTypeChecker
     m = prepare_pt2e(graph_module, quantizer)
@@ -206,9 +196,7 @@ def test_quantizer_conv2d_relu():
 
     example_input = (torch.ones(1, 4, 32, 32),)
     quantizer = NeutronQuantizer()
-    graph_module = torch.export.export_for_training(
-        model, example_input, strict=True
-    ).module()
+    graph_module = torch.export.export(model, example_input, strict=True).module()
 
     # noinspection PyTypeChecker
     m = prepare_pt2e(graph_module, quantizer)
@@ -231,9 +219,7 @@ def test_quantizer_conv2d_avg_pool2d():
 
     example_input = (torch.ones(1, 4, 16, 16),)
     quantizer = NeutronQuantizer()
-    graph_module = torch.export.export_for_training(
-        model, example_input, strict=True
-    ).module()
+    graph_module = torch.export.export(model, example_input, strict=True).module()
 
     # noinspection PyTypeChecker
     m = prepare_pt2e(graph_module, quantizer)
@@ -256,9 +242,7 @@ def test_quantizer_conv2d_permute():
 
     example_input = (torch.ones(1, 4, 16, 16),)
     quantizer = NeutronQuantizer()
-    graph_module = torch.export.export_for_training(
-        model, example_input, strict=True
-    ).module()
+    graph_module = torch.export.export(model, example_input, strict=True).module()
 
     # noinspection PyTypeChecker
     m = prepare_pt2e(graph_module, quantizer)
@@ -285,9 +269,7 @@ def test_multiple_shared_spec_ops_in_row():
 
     example_input = (torch.ones(1, 3, 64, 64),)
     quantizer = NeutronQuantizer()
-    graph_module = torch.export.export_for_training(
-        model, example_input, strict=True
-    ).module()
+    graph_module = torch.export.export(model, example_input, strict=True).module()
 
     # noinspection PyTypeChecker
     m = prepare_pt2e(graph_module, quantizer)
@@ -321,9 +303,7 @@ def test_quantizers_order_invariance():
     example_input = (torch.ones(1, 4, 64, 64),)
     quantizer = NeutronQuantizer()
 
-    graph_module = torch.export.export_for_training(
-        model, example_input, strict=True
-    ).module()
+    graph_module = torch.export.export(model, example_input, strict=True).module()
 
     m = prepare_pt2e(deepcopy(graph_module), quantizer)
     m(*example_input)

examples/nxp/aot_neutron_compile.py

@@ -15,7 +15,6 @@
 import executorch.kernels.quantized  # noqa F401
 
 import torch
-
 from executorch.backends.nxp.backend.ir.edge_passes.remove_io_quant_ops_pass import (
     RemoveIOQuantOpsPass,
 )
@@ -24,14 +23,12 @@
 from executorch.backends.nxp.quantizer.neutron_quantizer import NeutronQuantizer
 from executorch.examples.models import MODEL_NAME_TO_MODEL
 from executorch.examples.models.model_factory import EagerModelFactory
-
 from executorch.exir import (
     EdgeCompileConfig,
     ExecutorchBackendConfig,
     to_edge_transform_and_lower,
 )
 from executorch.extension.export_util import save_pte_program
-
 from torch.export import export
 from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e
 
@@ -227,9 +224,7 @@ def _get_batch_size(data):
     model = model.eval()
 
     # 2. Export the model to ATEN
-    exported_program = torch.export.export_for_training(
-        model, example_inputs, strict=True
-    )
+    exported_program = torch.export.export(model, example_inputs, strict=True)
 
     module = exported_program.module()