Explain what unlift config does

Summary: "Unlift" config depends on the "enable_aot" config. Let's make sure it is mentioned in the code.

Reviewed By: gmagogsfm

Differential Revision: D48181418

fbshipit-source-id: 91f83544f8daaf3814fdd6c4ed22785d571b3086

Author: tugsbayasgalan

exir/capture/_capture.py

@@ -60,8 +60,19 @@ def capture(
     out_spec = None
     # TODO (zhxchen17) Always functionalize in a second pass no matter which path is taken.
     flat_args = tuple(pytree.tree_flatten(args)[0])
+    if not config.enable_aot:
+        if config._unlift:
+            raise ExportError(
+                ExportErrorType.NOT_SUPPORTED,
+                "_unlift config doesn't do anything without enable_aot enabled. Please do not set it",
+            )
     if config.pt2_mode:
         if config.enable_aot:
+            if config.enable_dynamic_shape:
+                raise ExportError(
+                    ExportErrorType.NOT_SUPPORTED,
+                    "Under enable_aot, enable_dynamic_shapes flag doesn't do anything. Please do not set it",
+                )
             if not config.enable_functionalization:
                 raise ExportError(
                     ExportErrorType.NOT_SUPPORTED,

exir/capture/_config.py

@@ -22,7 +22,7 @@ class CaptureConfig:
     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
+    _unlift: bool = False  # This flag does nothing if enable_aot is False.
     _use_old_decomp_table: bool = False
 
 

exir/emit/test/test_emit.py

@@ -568,7 +568,6 @@ def test_model(eager_module):
             capture_config = exir.CaptureConfig(
                 pt2_mode=True,
                 enable_functionalization=True,
-                enable_dynamic_shape=True,
                 enable_aot=True,
                 _unlift=False,
             )