delete exir.capture from some more places (#2852)

Summary:

a bunch of misc places using exir.capture that shouldnt be anymore.

Reviewed By: mergennachin

Differential Revision: D55766467

Author: JacobSzwejbka

exir/tests/test_quantization.py

@@ -10,8 +10,7 @@
 
 import torch
 import torchvision
-from executorch import exir
-from executorch.exir import EdgeCompileConfig
+from executorch.exir import EdgeCompileConfig, to_edge
 from executorch.exir.passes.quant_fusion_pass import QuantFusionPass
 from executorch.exir.passes.spec_prop_pass import SpecPropPass
 from torch.ao.ns.fx.utils import compute_sqnr
@@ -29,6 +28,7 @@
     get_symmetric_quantization_config,
     XNNPACKQuantizer,
 )
+from torch.export import export
 from torch.testing import FileCheck
 from torch.testing._internal.common_quantized import override_quantized_engine
 
@@ -70,19 +70,17 @@ def test_resnet(self) -> None:
             compile_config = EdgeCompileConfig(
                 _check_ir_validity=False,
             )
-            m = (
-                exir.capture(m, example_inputs)
-                .to_edge(config=compile_config)
-                .transform(QuantFusionPass(), SpecPropPass())
-            )
+            m = to_edge(
+                export(m, example_inputs), compile_config=compile_config
+            ).transform([QuantFusionPass(), SpecPropPass()])
 
-            after_quant_result = m(*example_inputs)[0]
+            after_quant_result = m.exported_program().module()(*example_inputs)[0]
             FileCheck().check(
                 "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor"
             ).check(
                 "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor"
             ).run(
-                m.exported_program.graph_module.code
+                m.exported_program().graph_module.code
             )
             # after_quant_fusion_result = m(*example_inputs)[0]
 

test/models/generate_linear_out_bundled_program.py

@@ -16,10 +16,8 @@
 import subprocess
 from typing import List
 
-import executorch.exir as exir
-
 import torch
-from executorch.exir import ExecutorchBackendConfig
+from executorch.exir import ExecutorchBackendConfig, to_edge
 
 from executorch.exir.passes import MemoryPlanningPass, ToOutVarPass
 from executorch.exir.print_program import pretty_print
@@ -30,6 +28,7 @@
 )
 
 from executorch.test.models.linear_model import LinearModel
+from torch.export import export
 
 
 def main() -> None:
@@ -38,21 +37,17 @@ def main() -> None:
     trace_inputs = (torch.ones(2, 2, dtype=torch.float),)
 
     # Trace to FX Graph.
-    exec_prog = (
-        exir.capture(model, trace_inputs)
-        .to_edge()
-        .to_executorch(
-            config=ExecutorchBackendConfig(
-                memory_planning_pass=MemoryPlanningPass(),
-                to_out_var_pass=ToOutVarPass(),
-            )
+    exec_prog = to_edge(export(model, trace_inputs)).to_executorch(
+        config=ExecutorchBackendConfig(
+            memory_planning_pass=MemoryPlanningPass(),
+            to_out_var_pass=ToOutVarPass(),
         )
     )
     # Emit in-memory representation.
-    pretty_print(exec_prog.program)
+    pretty_print(exec_prog.executorch_program)
 
     # Serialize to flatbuffer.
-    exec_prog.program.version = 0
+    exec_prog.executorch_program.version = 0
 
     # Create test sets
     method_test_cases: List[MethodTestCase] = []
@@ -66,7 +61,6 @@ def main() -> None:
     ]
 
     bundled_program = BundledProgram(exec_prog, method_test_suites)
-    pretty_print(bundled_program)
 
     bundled_program_flatbuffer = serialize_from_bundled_program_to_flatbuffer(
         bundled_program
@@ -76,7 +70,7 @@ def main() -> None:
         subprocess.run(["hg", "root"], stdout=subprocess.PIPE).stdout.decode().strip()
     )
     with open(
-        f"{fbsource_base_path}/fbcode/executorch/test/models/linear_out_bundled_program.pte",
+        f"{fbsource_base_path}/fbcode/executorch/test/models/linear_out_bundled_program.bpte",
         "wb",
     ) as file:
         file.write(bundled_program_flatbuffer)