From 78cbac12d031a050d2bfc02600b210a5de52e611 Mon Sep 17 00:00:00 2001
From: Angela Yi <angelayi@meta.com>
Date: Mon, 22 Jan 2024 10:40:41 -0800
Subject: [PATCH] [export] Remove CallSpec (#117671)

Summary:

X-link: https://github.com/pytorch/executorch/pull/1618

This is not really being used anywhere

Test Plan: CI

Reviewed By: zhxchen17

Differential Revision: D52842563
---
 torch/_export/__init__.py         |  3 ---
 torch/_export/exported_program.py | 10 ----------
 torch/export/exported_program.py  |  3 ++-
 3 files changed, 2 insertions(+), 14 deletions(-)

diff --git a/torch/_export/__init__.py b/torch/_export/__init__.py
index da38079161fd..fcc27ba4d9e7 100644
--- a/torch/_export/__init__.py
+++ b/torch/_export/__init__.py
@@ -73,9 +73,6 @@
 from torch.fx.graph import _PyTreeCodeGen, _PyTreeInfo
 from torch.utils._sympy.value_ranges import ValueRangeError, ValueRanges
 
-from .exported_program import (
-    CallSpec,
-)
 from .passes.add_runtime_assertions_for_constraints_pass import (
     _AddRuntimeAssertionsForInlineConstraintsPass,
 )
diff --git a/torch/_export/exported_program.py b/torch/_export/exported_program.py
index a7b0fd679a5c..5d28ea315490 100644
--- a/torch/_export/exported_program.py
+++ b/torch/_export/exported_program.py
@@ -1,11 +1,8 @@
-import dataclasses
-from typing import Optional
 import warnings
 
 
 import torch
 import torch.fx
-import torch.utils._pytree as pytree
 
 
 # TODO(ycao): This is added to avoid breaking existing code temporarily.
@@ -32,13 +29,6 @@
 ]
 
 
-# Information to maintain user calling/returning specs
-@dataclasses.dataclass
-class CallSpec:
-    in_spec: Optional[pytree.TreeSpec]
-    out_spec: Optional[pytree.TreeSpec]
-
-
 def _create_graph_module_for_export(root, graph):
     try:
         gm = torch.fx.GraphModule(root, graph)
diff --git a/torch/export/exported_program.py b/torch/export/exported_program.py
index fe94e695b9cd..ae6063c3e2fa 100644
--- a/torch/export/exported_program.py
+++ b/torch/export/exported_program.py
@@ -1,4 +1,5 @@
 import copy
+from collections import namedtuple
 import dataclasses
 import functools
 import types
@@ -213,7 +214,7 @@ def example_inputs(self):
     @property
     @compatibility(is_backward_compatible=False)
     def call_spec(self):
-        from torch._export.exported_program import CallSpec
+        CallSpec = namedtuple("CallSpec", ["in_spec", "out_spec"])
 
         if len(self.module_call_graph) == 0:
             return CallSpec(in_spec=None, out_spec=None)