Fix export config in examples (#34)

Summary:
Pull Request resolved: #34

It doesn't look correct that want to export models with one config but test the export with a different one. Certain flags needs more clarification/documentation

This diff is to ensure using canonical config for both export and tests.

TODO: Rerun the existing tests failed with msg about missing required positional arguments. It's likely related to the params lifting and unlifting. Will figure it out.

Differential Revision: D48018569

fbshipit-source-id: fe45e72c0f406715d4c3a68e4dd1434a4d9caf58

Author: guangy10

examples/export/test/test_export.py

@@ -8,7 +8,7 @@
 
 import torch
 
-from executorch.examples.export.utils import _EDGE_COMPILE_CONFIG
+from executorch.examples.export.utils import _CAPTURE_CONFIG, _EDGE_COMPILE_CONFIG
 from executorch.examples.models import MODEL_NAME_TO_MODEL
 
 
@@ -18,16 +18,15 @@ def _assert_eager_lowered_same_result(
     ):
         import executorch.exir as exir
 
-        capture_config = exir.CaptureConfig(enable_dynamic_shape=False)
-        edge_model = exir.capture(eager_model, example_inputs, capture_config).to_edge(
+        edge_model = exir.capture(eager_model, example_inputs, _CAPTURE_CONFIG).to_edge(
             _EDGE_COMPILE_CONFIG
         )
 
         executorch_model = edge_model.to_executorch()
         with torch.no_grad():
             eager_output = eager_model(*example_inputs)
         with torch.no_grad():
-            executorch_output = executorch_model.graph_module(*example_inputs)
+            executorch_output = executorch_model(*example_inputs)
         self.assertTrue(
             torch.allclose(eager_output, executorch_output[0], rtol=1e-5, atol=1e-5)
         )

examples/export/utils.py

@@ -11,7 +11,7 @@
 # Reason is that there memory allocation ops with symbolic shape nodes.
 # and when evaulating shape, it doesnt seem that we presenting them with shape env
 # that contain those variables.
-_CAPTURE_CONFIG = exir.CaptureConfig(enable_aot=True, _unlift=False)
+_CAPTURE_CONFIG = exir.CaptureConfig(enable_aot=True)
 _EDGE_COMPILE_CONFIG = exir.EdgeCompileConfig(
     _check_ir_validity=False,
 )

exir/capture/_config.py

@@ -19,8 +19,8 @@
 class CaptureConfig:
     pt2_mode: bool = True
     enable_functionalization: bool = True
-    enable_dynamic_shape: bool = False
-    enable_aot: bool = False
+    enable_dynamic_shape: bool = False  # This flag does nothing if enable_aot is True
+    enable_aot: bool = False  # When it's true it implies automatic dynamic shapes via default dynamo config
     _dynamo_config: "ExirDynamoConfig" = field(default_factory=ExirDynamoConfig)
     _unlift: bool = False
     _use_old_decomp_table: bool = False

exir/program/_program.py

@@ -129,6 +129,9 @@ def __init__(
         self._constant_tensor_alignment: Optional[int] = constant_tensor_alignment
         self._delegate_alignment: Optional[int] = delegate_alignment
 
+    def __call__(self, *args: Any) -> Any:
+        return self.exported_program(*args)
+
     @property
     def buffer(self) -> bytes:
         if self._buffer is None: