Update export api to long term version

Summary: Update the export apis to long term version apis in bundled program code base and its documentations.

Differential Revision: D50256862

Author: Songhao Jia

bundled_program/tests/common.py

@@ -9,10 +9,12 @@
 import string
 from typing import List, Tuple, Union
 
-import executorch.exir as exir
 import torch
+
 from executorch.bundled_program.config import BundledConfig
+from executorch.exir import to_edge
 from executorch.exir.schema import Program
+from torch.export import export
 
 # @manual=fbsource//third-party/pypi/typing-extensions:typing-extensions
 from typing_extensions import TypeAlias
@@ -210,12 +212,14 @@ def get_common_program() -> Tuple[Program, BundledConfig]:
         for m_name in eager_model.method_names
     }
 
-    program = (
-        exir.capture_multiple(eager_model, capture_inputs)
-        .to_edge()
-        .to_executorch()
-        .program
-    )
+    # Trace to FX Graph and emit the program
+    method_graphs = {
+        m_name: export(getattr(eager_model, m_name), capture_inputs[m_name])
+        for m_name in eager_model.method_names
+    }
+
+    program = to_edge(method_graphs).to_executorch().executorch_program
+
     _, bundled_config = get_random_config_with_eager_model(
         eager_model=eager_model,
         method_names=eager_model.method_names,

docs/source/sdk-bundled-io.md

@@ -67,15 +67,13 @@ Here is a flow highlighting how to generate a `BundledProgram` given a PyTorch m
 ```python
 
 import torch
-from executorch import exir
+from torch.export import export
+
 from executorch.bundled_program.config import BundledConfig
 from executorch.bundled_program.core import create_bundled_program
 from executorch.bundled_program.serialize import serialize_from_bundled_program_to_flatbuffer
-from executorch.bundled_program.serialize import deserialize_from_flatbuffer_to_bundled_program
 
-
-from executorch.exir import ExecutorchBackendConfig
-from executorch.exir.passes import MemoryPlanningPass, ToOutVarPass
+from executorch.exir import to_edge
 
 
 # Step 1: ExecuTorch Program Export
@@ -119,12 +117,12 @@ capture_inputs = {
 }
 
 # Trace to FX Graph and emit the program
-program = (
-    exir.capture_multiple(model, capture_inputs)
-    .to_edge()
-    .to_executorch()
-    .program
-)
+method_graphs = {
+    m_name: export(getattr(model, m_name), capture_inputs[m_name])
+    for m_name in method_names
+}
+
+program = to_edge(method_graphs).to_executorch().executorch_program
 
 # Step 2: Construct BundledConfig
 
@@ -291,12 +289,13 @@ Here's the example of the dtype of test input not meet model's requirement:
 
 ```python
 import torch
-from executorch import exir
-from executorch.exir import ExecutorchBackendConfig
-from executorch.exir.passes import MemoryPlanningPass, ToOutVarPass
+from torch.export import export
+
 from executorch.bundled_program.config import BundledConfig
 from executorch.bundled_program.core import create_bundled_program
 
+from executorch.exir import to_edge
+
 
 class Module(torch.nn.Module):
     def __init__(self):
@@ -318,16 +317,13 @@ method_names = ['forward']
 inputs = torch.ones(2, 2, dtype=torch.float)
 print(model(inputs))
 
-# Trace to FX Graph.
-program = (
-    exir.capture(model, (inputs,))
-    .to_edge()
-    .to_executorch(
-        config=ExecutorchBackendConfig(
-            memory_planning_pass=MemoryPlanningPass(), to_out_var_pass=ToOutVarPass()
-        )
-    ).program
-)
+# Trace to FX Graph and emit the program
+method_graphs = {
+    m_name: export(getattr(model, m_name), (inputs, ))
+    for m_name in method_names
+}
+
+program = to_edge(method_graphs).to_executorch().executorch_program
 
 
 # number of input sets to be verified
@@ -416,12 +412,13 @@ Another common error would be the method name in `BundledConfig` does not exist
 
 ```python
 import torch
-from executorch import exir
-from executorch.exir import ExecutorchBackendConfig
-from executorch.exir.passes import MemoryPlanningPass, ToOutVarPass
+from torch.export import export
+
 from executorch.bundled_program.config import BundledConfig
 from executorch.bundled_program.core import create_bundled_program
 
+from executorch.exir import to_edge
+
 
 
 class Module(torch.nn.Module):
@@ -440,23 +437,17 @@ class Module(torch.nn.Module):
 
 model = Module()
 
-# NOTE: wrong_forward is not an inference method in the above model.
-method_names = ['wrong_forward']
+method_names = ['forward']
 
 inputs = torch.ones(2, 2, dtype=torch.float)
-print(model(inputs))
 
-# Trace to FX Graph.
-program = (
-    exir.capture(model, (inputs,))
-    .to_edge()
-    .to_executorch(
-        config=ExecutorchBackendConfig(
-            memory_planning_pass=MemoryPlanningPass(), to_out_var_pass=ToOutVarPass()
-        )
-    ).program
-)
+# Trace to FX Graph and emit the program
+method_graphs = {
+    m_name: export(getattr(model, m_name), (inputs, ))
+    for m_name in method_names
+}
 
+program = to_edge(method_graphs).to_executorch().executorch_program
 
 # Number of input sets to be verified
 n_input = 10
@@ -476,7 +467,11 @@ expected_outpus = [
     [[model(*x)] for x in inputs[0]]
 ]
 
-bundled_config = BundledConfig(method_names, inputs, expected_outpus)
+
+# NOTE: WRONG_FORWARD is not an inference method in the above model.
+wrong_method_names = ['WRONG_FORWARD']
+
+bundled_config = BundledConfig(wrong_method_names, inputs, expected_outpus)
 
 bundled_program = create_bundled_program(program, bundled_config)