Remove CallSpec (pytorch#1618)

Summary:
X-link: pytorch/pytorch#117671


This is not really being used anywhere

Differential Revision: D52842563

exir/__init__.py

@@ -28,8 +28,7 @@
     to_edge,
 )
 from executorch.exir.tracer import ExirDynamoConfig
-from torch._export import (  # lots of people are doing from exir import CallSpec, ExportGraphSignature, ExportedProgram which seems wrong
-    CallSpec,
+from torch.export import (
     ExportedProgram,
     ExportGraphSignature,
 )
@@ -42,7 +41,6 @@
     "capture",
     "capture_multiple",
     "_capture_legacy_do_not_use",
-    "CallSpec",
     "ExportedProgram",
     "ExirExportedProgram",
     "ExecutorchProgram",

exir/capture/_capture.py

@@ -5,6 +5,7 @@
 # LICENSE file in the root directory of this source tree.
 
 import copy
+import dataclasses
 import warnings
 from collections import namedtuple
 from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
@@ -26,10 +27,12 @@
 from torch import _guards
 from torch._dispatch.python import enable_python_dispatcher
 from torch._dynamo.eval_frame import Constraint
-from torch._export import CallSpec, export, ExportedProgram, ExportGraphSignature
 from torch._export.passes import ReplaceViewOpsWithViewCopyOpsPass
 from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode
+from torch.export import export
 from torch.export.exported_program import (
+    ExportedProgram,
+    ExportGraphSignature,
     InputKind,
     InputSpec,
     ModuleCallEntry,
@@ -52,6 +55,11 @@
     "CompileSpec", ["method_name", "callable", "args", "constraints"]
 )
 
+@dataclasses.dataclass
+class CallSpec:
+    in_spec: Optional[pytree.TreeSpec]
+    out_spec: Optional[pytree.TreeSpec]
+
 
 @compatibility(is_backward_compatible=False)
 def _capture_legacy_do_not_use(f, args) -> ExirExportedProgram:

exir/lowered_backend_module.py

@@ -271,10 +271,6 @@ def program(self, emit_stacktrace: bool = False) -> Program:
                 ),
                 lowered_exported_program.graph_module,
             ),
-            # TODO: May need to set lowered_exported_program.call_spec = CallSpec(None, None)
-            # somewhere as we should pass it a list of tensors to the lowered module and output a
-            # list of tensors. Putting call_spec=lowered_exported_program.call_spec is correct here as the
-            # inputs/outputs to the toplevel program will be in the format of the eager module.
             state_dict={},  # None because all data are consumed by delegate
             range_constraints=lowered_exported_program.range_constraints,
             module_call_graph=lowered_exported_program.module_call_graph,