We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
✨[Feature] Create Decorator for Python functions Incompatible with Legacy Torch Versions
## Context
Some functions in the code base are incompatible with certain versions of Torch such as 1.13.1. Adding a general-purpose utility decorator which detects whether the Torch version is incompatible with certain key functions could be helpful in notifying the user when such mismatches occur.

## Desired Solution
A decorator of the form:
```python
@req_torch_version(2)
def f(...):
```
This decorator would throw a clear error when the detected Torch version differs from that required to use the decorated function.

## Temporary Alternative
A temporary alternative, which is already in use in PR #1731, is to add logic to each function with compatibility issues, raising an error if any versioning issues are detected. While this is a functional solution, it is difficult to scale and can lead to repetitive code.

## Additional Context
Certain imports are also not compatible across Torch versions (as in `import torchdynamo` in Torch 1.13 versus `import torch._dynamo` in 2.0). To resolve this issue, the imports can be moved inside the functions using them so as to encapsulate all version-specific code within one area.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `py/torch_tensorrt/fx/tracer/dispatch_tracer/aten_tracer.py`
Content:
```
1 import copy
2 import sys
3 from contextlib import contextmanager
4 from typing import Any, Callable, Dict, Generator, List, Optional, Set, Tuple, Union
5 
6 import torch
7 
8 if not torch.__version__.startswith("1"):
9     import torch._dynamo as torchdynamo
10 
11 from torch.fx.passes.infra.pass_base import PassResult
12 
13 from torch_tensorrt.fx.passes.lower_basic_pass_aten import (
14     compose_bmm,
15     compose_chunk,
16     compose_getitem_slice,
17     remove_ops,
18     replace_aten_op_with_indices,
19     replace_aten_reshape_alias_with_replace,
20     replace_builtin_ops,
21     replace_inplace_ops,
22     replace_native_layernorm_with_layernorm,
23     replace_transpose_mm_op_with_linear,
24     run_const_fold,
25 )
26 from typing_extensions import TypeAlias
27 
28 Value: TypeAlias = Union[
29     Tuple["Value", ...],
30     List["Value"],
31     Dict[str, "Value"],
32 ]
33 
34 
35 class DynamoConfig:
36     """
37     Manage Exir-specific configurations of Dynamo.
38     """
39 
40     def __init__(
41         self,
42         capture_scalar_outputs: bool = True,
43         guard_nn_modules: bool = True,
44         dynamic_shapes: bool = True,
45         specialize_int_float: bool = True,
46         verbose: bool = True,
47     ) -> None:
48 
49         self.capture_scalar_outputs = capture_scalar_outputs
50         self.guard_nn_modules = guard_nn_modules
51         self.dynamic_shapes = dynamic_shapes
52         self.specialize_int_float = specialize_int_float
53         self.verbose = verbose
54 
55     def activate(self) -> None:
56         torchdynamo.config.capture_scalar_outputs = self.capture_scalar_outputs
57         torchdynamo.config.guard_nn_modules = self.guard_nn_modules
58         torchdynamo.config.dynamic_shapes = self.dynamic_shapes
59         torchdynamo.config.specialize_int_float = self.specialize_int_float
60         torchdynamo.config.verbose = self.verbose
61 
62     def deactivate(self) -> None:
63         torchdynamo.config.capture_scalar_outputs = True
64         torchdynamo.config.guard_nn_modules = True
65         torchdynamo.config.dynamic_shapes = True
66         torchdynamo.config.specialize_int_float = True
67         torchdynamo.config.verbose = True
68 
69 
70 @contextmanager
71 def using_config(config: DynamoConfig) -> Generator[DynamoConfig, None, None]:
72     config.activate()
73     try:
74         yield config
75     finally:
76         config.deactivate()
77 
78 
79 @contextmanager
80 def setting_python_recursive_limit(limit: int = 10000) -> Generator[None, None, None]:
81     """
82     Temporarily increase the python interpreter stack recursion limit.
83     This is mostly used for pickling large scale modules.
84     """
85     default = sys.getrecursionlimit()
86     if limit > default:
87         sys.setrecursionlimit(limit)
88     try:
89         yield
90     finally:
91         sys.setrecursionlimit(default)
92 
93 
94 def dynamo_trace(
95     f: Callable[..., Value],
96     # pyre-ignore
97     args: Tuple[Any, ...],
98     aten_graph: bool,
99     tracing_mode: str = "real",
100     dynamo_config: Optional[DynamoConfig] = None,
101 ) -> Tuple[torch.fx.GraphModule, Set]:
102     """
103     TODO: Once we fully migrate to torchdynamo frontend, we will remove
104     this config option alltogether.  For now, it helps with quick
105     experiments with playing around with TorchDynamo
106     """
107     if torch.__version__.startswith("1"):
108         raise ValueError(
109             f"The aten tracer requires Torch version >= 2.0. Detected version {torch.__version__}"
110         )
111 
112     if dynamo_config is None:
113         dynamo_config = DynamoConfig()
114     with using_config(dynamo_config), setting_python_recursive_limit(2000):
115         torchdynamo.reset()
116         try:
117             return torchdynamo.export(
118                 f,
119                 *copy.deepcopy(args),
120                 aten_graph=aten_graph,
121                 tracing_mode=tracing_mode,
122             )
123         except torchdynamo.exc.Unsupported as exc:
124             raise RuntimeError(
125                 "The user code is using a feature we don't support. "
126                 "Please try torchdynamo.explain() to get possible the reasons",
127             ) from exc
128         except Exception as exc:
129             raise RuntimeError(
130                 "torchdynamo internal error occured. Please see above stacktrace"
131             ) from exc
132 
133 
134 def trace(f, args, *rest):
135     graph_module, guards = dynamo_trace(f, args, True, "symbolic")
136     return graph_module, guards
137 
138 
139 def opt_trace(f, args, *rest):
140     """
141     Optimized trace with necessary passes which re-compose some ops or replace some ops
142     These passes should be general and functional purpose
143     """
144     passes_list = [
145         compose_bmm,
146         compose_chunk,
147         compose_getitem_slice,
148         replace_aten_reshape_alias_with_replace,
149         replace_aten_op_with_indices,
150         replace_transpose_mm_op_with_linear,  # after compose_bmm
151         replace_native_layernorm_with_layernorm,
152         remove_ops,
153         replace_builtin_ops,  # after replace_native_layernorm_with_layernorm
154         replace_inplace_ops,  # remove it once functionalization is enabled
155     ]
156 
157     fx_module, _ = trace(f, args)
158     print(fx_module.graph)
159     for passes in passes_list:
160         pr: PassResult = passes(fx_module)
161         fx_module = pr.graph_module
162 
163     fx_module(*args)
164 
165     fx_module = run_const_fold(fx_module)
166     print(fx_module.graph)
167     return fx_module
168
```

Path: `py/torch_tensorrt/fx/utils.py`
Content:
```
1 from enum import Enum
2 from typing import List
3 
4 # @manual=//deeplearning/trt/python:py_tensorrt
5 import tensorrt as trt
6 import torch
7 from functorch import make_fx
8 from functorch.experimental import functionalize
9 from torch_tensorrt.fx.passes.lower_basic_pass import (
10     replace_op_with_indices,
11     run_const_fold,
12 )
13 
14 from .types import Shape, TRTDataType
15 
16 
17 class LowerPrecision(Enum):
18     FP32 = "fp32"
19     FP16 = "fp16"
20     INT8 = "int8"
21 
22 
23 def torch_dtype_to_trt(dtype: torch.dtype) -> TRTDataType:
24     """
25     Convert PyTorch data types to TensorRT data types.
26 
27     Args:
28         dtype (torch.dtype): A PyTorch data type.
29 
30     Returns:
31         The equivalent TensorRT data type.
32     """
33     if trt.__version__ >= "7.0" and dtype == torch.bool:
34         return trt.bool
35     elif dtype == torch.int8:
36         return trt.int8
37     elif dtype == torch.int32:
38         return trt.int32
39     elif dtype == torch.float16:
40         return trt.float16
41     elif dtype == torch.float32:
42         return trt.float32
43     else:
44         raise TypeError("%s is not supported by tensorrt" % dtype)
45 
46 
47 def torch_dtype_from_trt(dtype: TRTDataType) -> torch.dtype:
48     """
49     Convert TensorRT data types to PyTorch data types.
50 
51     Args:
52         dtype (TRTDataType): A TensorRT data type.
53 
54     Returns:
55         The equivalent PyTorch data type.
56     """
57     if dtype == trt.int8:
58         return torch.int8
59     elif trt.__version__ >= "7.0" and dtype == trt.bool:
60         return torch.bool
61     elif dtype == trt.int32:
62         return torch.int32
63     elif dtype == trt.float16:
64         return torch.float16
65     elif dtype == trt.float32:
66         return torch.float32
67     else:
68         raise TypeError("%s is not supported by torch" % dtype)
69 
70 
71 def get_dynamic_dims(shape: Shape) -> List[int]:
72     """
73     This function finds the dynamic dimensions in the given
74     shape. A dimension is dynamic if it's -1.
75 
76     Args:
77         shape (Shape): A sequence of integer that represents
78             the shape of a tensor.
79 
80     Returns:
81         A list of integers contains all the dynamic dimensions
82         in the given shape
83     """
84     dynamic_dims = []
85 
86     for i, s in enumerate(shape):
87         if s == -1:
88             dynamic_dims.append(i)
89 
90     return dynamic_dims
91 
92 
93 def proxytensor_trace(mod, inputs):
94 
95     mod.eval()
96 
97     def f(*inp):
98         return mod(*inp)
99 
100     mod = make_fx(functionalize(f))(*inputs)
101 
102     # Remove const operation. For ex, nn.Linear has transpose operation on weight
103     mod.graph.eliminate_dead_code()
104     mod = run_const_fold(mod)
105     mod = replace_op_with_indices(mod)
106     return mod
107
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/py/torch_tensorrt/fx/tracer/dispatch_tracer/aten_tracer.py b/py/torch_tensorrt/fx/tracer/dispatch_tracer/aten_tracer.py
--- a/py/torch_tensorrt/fx/tracer/dispatch_tracer/aten_tracer.py
+++ b/py/torch_tensorrt/fx/tracer/dispatch_tracer/aten_tracer.py
@@ -9,7 +9,7 @@
     import torch._dynamo as torchdynamo
 
 from torch.fx.passes.infra.pass_base import PassResult
-
+from torch_tensorrt.fx.utils import req_torch_version
 from torch_tensorrt.fx.passes.lower_basic_pass_aten import (
     compose_bmm,
     compose_chunk,
@@ -91,6 +91,7 @@
         sys.setrecursionlimit(default)
 
 
+@req_torch_version("2.0")
 def dynamo_trace(
     f: Callable[..., Value],
     # pyre-ignore
@@ -104,11 +105,6 @@
     this config option alltogether.  For now, it helps with quick
     experiments with playing around with TorchDynamo
     """
-    if torch.__version__.startswith("1"):
-        raise ValueError(
-            f"The aten tracer requires Torch version >= 2.0. Detected version {torch.__version__}"
-        )
-
     if dynamo_config is None:
         dynamo_config = DynamoConfig()
     with using_config(dynamo_config), setting_python_recursive_limit(2000):
@@ -131,11 +127,13 @@
             ) from exc
 
 
+@req_torch_version("2.0")
 def trace(f, args, *rest):
     graph_module, guards = dynamo_trace(f, args, True, "symbolic")
     return graph_module, guards
 
 
+@req_torch_version("2.0")
 def opt_trace(f, args, *rest):
     """
     Optimized trace with necessary passes which re-compose some ops or replace some ops
diff --git a/py/torch_tensorrt/fx/utils.py b/py/torch_tensorrt/fx/utils.py
--- a/py/torch_tensorrt/fx/utils.py
+++ b/py/torch_tensorrt/fx/utils.py
@@ -1,5 +1,6 @@
 from enum import Enum
-from typing import List
+from typing import List, Callable
+from packaging import version
 
 # @manual=//deeplearning/trt/python:py_tensorrt
 import tensorrt as trt
@@ -104,3 +105,36 @@
     mod = run_const_fold(mod)
     mod = replace_op_with_indices(mod)
     return mod
+
+
+def req_torch_version(min_torch_version: str = "2.dev"):
+    """
+    Create a decorator which verifies the Torch version installed
+    against a specified version range
+
+    Args:
+        min_torch_version (str): The minimum required Torch version
+        for the decorated function to work properly
+
+    Returns:
+        A decorator which raises a descriptive error message if
+        an unsupported Torch version is used
+    """
+
+    def nested_decorator(f: Callable):
+        def function_wrapper(*args, **kwargs):
+            # Parse minimum and current Torch versions
+            min_version = version.parse(min_torch_version)
+            current_version = version.parse(torch.__version__)
+
+            if current_version < min_version:
+                raise AssertionError(
+                    f"Expected Torch version {min_torch_version} or greater, "
+                    + f"when calling {f}. Detected version {torch.__version__}"
+                )
+            else:
+                return f(*args, **kwargs)
+
+        return function_wrapper
+
+    return nested_decorator

{"golden_diff": "diff --git a/py/torch_tensorrt/fx/tracer/dispatch_tracer/aten_tracer.py b/py/torch_tensorrt/fx/tracer/dispatch_tracer/aten_tracer.py\n--- a/py/torch_tensorrt/fx/tracer/dispatch_tracer/aten_tracer.py\n+++ b/py/torch_tensorrt/fx/tracer/dispatch_tracer/aten_tracer.py\n@@ -9,7 +9,7 @@\n     import torch._dynamo as torchdynamo\n \n from torch.fx.passes.infra.pass_base import PassResult\n-\n+from torch_tensorrt.fx.utils import req_torch_version\n from torch_tensorrt.fx.passes.lower_basic_pass_aten import (\n     compose_bmm,\n     compose_chunk,\n@@ -91,6 +91,7 @@\n         sys.setrecursionlimit(default)\n \n \n+@req_torch_version(\"2.0\")\n def dynamo_trace(\n     f: Callable[..., Value],\n     # pyre-ignore\n@@ -104,11 +105,6 @@\n     this config option alltogether.  For now, it helps with quick\n     experiments with playing around with TorchDynamo\n     \"\"\"\n-    if torch.__version__.startswith(\"1\"):\n-        raise ValueError(\n-            f\"The aten tracer requires Torch version >= 2.0. Detected version {torch.__version__}\"\n-        )\n-\n     if dynamo_config is None:\n         dynamo_config = DynamoConfig()\n     with using_config(dynamo_config), setting_python_recursive_limit(2000):\n@@ -131,11 +127,13 @@\n             ) from exc\n \n \n+@req_torch_version(\"2.0\")\n def trace(f, args, *rest):\n     graph_module, guards = dynamo_trace(f, args, True, \"symbolic\")\n     return graph_module, guards\n \n \n+@req_torch_version(\"2.0\")\n def opt_trace(f, args, *rest):\n     \"\"\"\n     Optimized trace with necessary passes which re-compose some ops or replace some ops\ndiff --git a/py/torch_tensorrt/fx/utils.py b/py/torch_tensorrt/fx/utils.py\n--- a/py/torch_tensorrt/fx/utils.py\n+++ b/py/torch_tensorrt/fx/utils.py\n@@ -1,5 +1,6 @@\n from enum import Enum\n-from typing import List\n+from typing import List, Callable\n+from packaging import version\n \n # @manual=//deeplearning/trt/python:py_tensorrt\n import tensorrt as trt\n@@ -104,3 +105,36 @@\n     mod = run_const_fold(mod)\n     mod = replace_op_with_indices(mod)\n     return mod\n+\n+\n+def req_torch_version(min_torch_version: str = \"2.dev\"):\n+    \"\"\"\n+    Create a decorator which verifies the Torch version installed\n+    against a specified version range\n+\n+    Args:\n+        min_torch_version (str): The minimum required Torch version\n+        for the decorated function to work properly\n+\n+    Returns:\n+        A decorator which raises a descriptive error message if\n+        an unsupported Torch version is used\n+    \"\"\"\n+\n+    def nested_decorator(f: Callable):\n+        def function_wrapper(*args, **kwargs):\n+            # Parse minimum and current Torch versions\n+            min_version = version.parse(min_torch_version)\n+            current_version = version.parse(torch.__version__)\n+\n+            if current_version < min_version:\n+                raise AssertionError(\n+                    f\"Expected Torch version {min_torch_version} or greater, \"\n+                    + f\"when calling {f}. Detected version {torch.__version__}\"\n+                )\n+            else:\n+                return f(*args, **kwargs)\n+\n+        return function_wrapper\n+\n+    return nested_decorator\n", "issue": "\u2728[Feature] Create Decorator for Python functions Incompatible with Legacy Torch Versions\n## Context\r\nSome functions in the code base are incompatible with certain versions of Torch such as 1.13.1. Adding a general-purpose utility decorator which detects whether the Torch version is incompatible with certain key functions could be helpful in notifying the user when such mismatches occur.\r\n\r\n## Desired Solution\r\nA decorator of the form:\r\n```python\r\n@req_torch_version(2)\r\ndef f(...):\r\n```\r\nThis decorator would throw a clear error when the detected Torch version differs from that required to use the decorated function.\r\n\r\n## Temporary Alternative\r\nA temporary alternative, which is already in use in PR #1731, is to add logic to each function with compatibility issues, raising an error if any versioning issues are detected. While this is a functional solution, it is difficult to scale and can lead to repetitive code.\r\n\r\n## Additional Context\r\nCertain imports are also not compatible across Torch versions (as in `import torchdynamo` in Torch 1.13 versus `import torch._dynamo` in 2.0). To resolve this issue, the imports can be moved inside the functions using them so as to encapsulate all version-specific code within one area.\n", "before_files": [{"content": "import copy\nimport sys\nfrom contextlib import contextmanager\nfrom typing import Any, Callable, Dict, Generator, List, Optional, Set, Tuple, Union\n\nimport torch\n\nif not torch.__version__.startswith(\"1\"):\n    import torch._dynamo as torchdynamo\n\nfrom torch.fx.passes.infra.pass_base import PassResult\n\nfrom torch_tensorrt.fx.passes.lower_basic_pass_aten import (\n    compose_bmm,\n    compose_chunk,\n    compose_getitem_slice,\n    remove_ops,\n    replace_aten_op_with_indices,\n    replace_aten_reshape_alias_with_replace,\n    replace_builtin_ops,\n    replace_inplace_ops,\n    replace_native_layernorm_with_layernorm,\n    replace_transpose_mm_op_with_linear,\n    run_const_fold,\n)\nfrom typing_extensions import TypeAlias\n\nValue: TypeAlias = Union[\n    Tuple[\"Value\", ...],\n    List[\"Value\"],\n    Dict[str, \"Value\"],\n]\n\n\nclass DynamoConfig:\n    \"\"\"\n    Manage Exir-specific configurations of Dynamo.\n    \"\"\"\n\n    def __init__(\n        self,\n        capture_scalar_outputs: bool = True,\n        guard_nn_modules: bool = True,\n        dynamic_shapes: bool = True,\n        specialize_int_float: bool = True,\n        verbose: bool = True,\n    ) -> None:\n\n        self.capture_scalar_outputs = capture_scalar_outputs\n        self.guard_nn_modules = guard_nn_modules\n        self.dynamic_shapes = dynamic_shapes\n        self.specialize_int_float = specialize_int_float\n        self.verbose = verbose\n\n    def activate(self) -> None:\n        torchdynamo.config.capture_scalar_outputs = self.capture_scalar_outputs\n        torchdynamo.config.guard_nn_modules = self.guard_nn_modules\n        torchdynamo.config.dynamic_shapes = self.dynamic_shapes\n        torchdynamo.config.specialize_int_float = self.specialize_int_float\n        torchdynamo.config.verbose = self.verbose\n\n    def deactivate(self) -> None:\n        torchdynamo.config.capture_scalar_outputs = True\n        torchdynamo.config.guard_nn_modules = True\n        torchdynamo.config.dynamic_shapes = True\n        torchdynamo.config.specialize_int_float = True\n        torchdynamo.config.verbose = True\n\n\n@contextmanager\ndef using_config(config: DynamoConfig) -> Generator[DynamoConfig, None, None]:\n    config.activate()\n    try:\n        yield config\n    finally:\n        config.deactivate()\n\n\n@contextmanager\ndef setting_python_recursive_limit(limit: int = 10000) -> Generator[None, None, None]:\n    \"\"\"\n    Temporarily increase the python interpreter stack recursion limit.\n    This is mostly used for pickling large scale modules.\n    \"\"\"\n    default = sys.getrecursionlimit()\n    if limit > default:\n        sys.setrecursionlimit(limit)\n    try:\n        yield\n    finally:\n        sys.setrecursionlimit(default)\n\n\ndef dynamo_trace(\n    f: Callable[..., Value],\n    # pyre-ignore\n    args: Tuple[Any, ...],\n    aten_graph: bool,\n    tracing_mode: str = \"real\",\n    dynamo_config: Optional[DynamoConfig] = None,\n) -> Tuple[torch.fx.GraphModule, Set]:\n    \"\"\"\n    TODO: Once we fully migrate to torchdynamo frontend, we will remove\n    this config option alltogether.  For now, it helps with quick\n    experiments with playing around with TorchDynamo\n    \"\"\"\n    if torch.__version__.startswith(\"1\"):\n        raise ValueError(\n            f\"The aten tracer requires Torch version >= 2.0. Detected version {torch.__version__}\"\n        )\n\n    if dynamo_config is None:\n        dynamo_config = DynamoConfig()\n    with using_config(dynamo_config), setting_python_recursive_limit(2000):\n        torchdynamo.reset()\n        try:\n            return torchdynamo.export(\n                f,\n                *copy.deepcopy(args),\n                aten_graph=aten_graph,\n                tracing_mode=tracing_mode,\n            )\n        except torchdynamo.exc.Unsupported as exc:\n            raise RuntimeError(\n                \"The user code is using a feature we don't support. \"\n                \"Please try torchdynamo.explain() to get possible the reasons\",\n            ) from exc\n        except Exception as exc:\n            raise RuntimeError(\n                \"torchdynamo internal error occured. Please see above stacktrace\"\n            ) from exc\n\n\ndef trace(f, args, *rest):\n    graph_module, guards = dynamo_trace(f, args, True, \"symbolic\")\n    return graph_module, guards\n\n\ndef opt_trace(f, args, *rest):\n    \"\"\"\n    Optimized trace with necessary passes which re-compose some ops or replace some ops\n    These passes should be general and functional purpose\n    \"\"\"\n    passes_list = [\n        compose_bmm,\n        compose_chunk,\n        compose_getitem_slice,\n        replace_aten_reshape_alias_with_replace,\n        replace_aten_op_with_indices,\n        replace_transpose_mm_op_with_linear,  # after compose_bmm\n        replace_native_layernorm_with_layernorm,\n        remove_ops,\n        replace_builtin_ops,  # after replace_native_layernorm_with_layernorm\n        replace_inplace_ops,  # remove it once functionalization is enabled\n    ]\n\n    fx_module, _ = trace(f, args)\n    print(fx_module.graph)\n    for passes in passes_list:\n        pr: PassResult = passes(fx_module)\n        fx_module = pr.graph_module\n\n    fx_module(*args)\n\n    fx_module = run_const_fold(fx_module)\n    print(fx_module.graph)\n    return fx_module\n", "path": "py/torch_tensorrt/fx/tracer/dispatch_tracer/aten_tracer.py"}, {"content": "from enum import Enum\nfrom typing import List\n\n# @manual=//deeplearning/trt/python:py_tensorrt\nimport tensorrt as trt\nimport torch\nfrom functorch import make_fx\nfrom functorch.experimental import functionalize\nfrom torch_tensorrt.fx.passes.lower_basic_pass import (\n    replace_op_with_indices,\n    run_const_fold,\n)\n\nfrom .types import Shape, TRTDataType\n\n\nclass LowerPrecision(Enum):\n    FP32 = \"fp32\"\n    FP16 = \"fp16\"\n    INT8 = \"int8\"\n\n\ndef torch_dtype_to_trt(dtype: torch.dtype) -> TRTDataType:\n    \"\"\"\n    Convert PyTorch data types to TensorRT data types.\n\n    Args:\n        dtype (torch.dtype): A PyTorch data type.\n\n    Returns:\n        The equivalent TensorRT data type.\n    \"\"\"\n    if trt.__version__ >= \"7.0\" and dtype == torch.bool:\n        return trt.bool\n    elif dtype == torch.int8:\n        return trt.int8\n    elif dtype == torch.int32:\n        return trt.int32\n    elif dtype == torch.float16:\n        return trt.float16\n    elif dtype == torch.float32:\n        return trt.float32\n    else:\n        raise TypeError(\"%s is not supported by tensorrt\" % dtype)\n\n\ndef torch_dtype_from_trt(dtype: TRTDataType) -> torch.dtype:\n    \"\"\"\n    Convert TensorRT data types to PyTorch data types.\n\n    Args:\n        dtype (TRTDataType): A TensorRT data type.\n\n    Returns:\n        The equivalent PyTorch data type.\n    \"\"\"\n    if dtype == trt.int8:\n        return torch.int8\n    elif trt.__version__ >= \"7.0\" and dtype == trt.bool:\n        return torch.bool\n    elif dtype == trt.int32:\n        return torch.int32\n    elif dtype == trt.float16:\n        return torch.float16\n    elif dtype == trt.float32:\n        return torch.float32\n    else:\n        raise TypeError(\"%s is not supported by torch\" % dtype)\n\n\ndef get_dynamic_dims(shape: Shape) -> List[int]:\n    \"\"\"\n    This function finds the dynamic dimensions in the given\n    shape. A dimension is dynamic if it's -1.\n\n    Args:\n        shape (Shape): A sequence of integer that represents\n            the shape of a tensor.\n\n    Returns:\n        A list of integers contains all the dynamic dimensions\n        in the given shape\n    \"\"\"\n    dynamic_dims = []\n\n    for i, s in enumerate(shape):\n        if s == -1:\n            dynamic_dims.append(i)\n\n    return dynamic_dims\n\n\ndef proxytensor_trace(mod, inputs):\n\n    mod.eval()\n\n    def f(*inp):\n        return mod(*inp)\n\n    mod = make_fx(functionalize(f))(*inputs)\n\n    # Remove const operation. For ex, nn.Linear has transpose operation on weight\n    mod.graph.eliminate_dead_code()\n    mod = run_const_fold(mod)\n    mod = replace_op_with_indices(mod)\n    return mod\n", "path": "py/torch_tensorrt/fx/utils.py"}], "after_files": [{"content": "import copy\nimport sys\nfrom contextlib import contextmanager\nfrom typing import Any, Callable, Dict, Generator, List, Optional, Set, Tuple, Union\n\nimport torch\n\nif not torch.__version__.startswith(\"1\"):\n    import torch._dynamo as torchdynamo\n\nfrom torch.fx.passes.infra.pass_base import PassResult\nfrom torch_tensorrt.fx.utils import req_torch_version\nfrom torch_tensorrt.fx.passes.lower_basic_pass_aten import (\n    compose_bmm,\n    compose_chunk,\n    compose_getitem_slice,\n    remove_ops,\n    replace_aten_op_with_indices,\n    replace_aten_reshape_alias_with_replace,\n    replace_builtin_ops,\n    replace_inplace_ops,\n    replace_native_layernorm_with_layernorm,\n    replace_transpose_mm_op_with_linear,\n    run_const_fold,\n)\nfrom typing_extensions import TypeAlias\n\nValue: TypeAlias = Union[\n    Tuple[\"Value\", ...],\n    List[\"Value\"],\n    Dict[str, \"Value\"],\n]\n\n\nclass DynamoConfig:\n    \"\"\"\n    Manage Exir-specific configurations of Dynamo.\n    \"\"\"\n\n    def __init__(\n        self,\n        capture_scalar_outputs: bool = True,\n        guard_nn_modules: bool = True,\n        dynamic_shapes: bool = True,\n        specialize_int_float: bool = True,\n        verbose: bool = True,\n    ) -> None:\n\n        self.capture_scalar_outputs = capture_scalar_outputs\n        self.guard_nn_modules = guard_nn_modules\n        self.dynamic_shapes = dynamic_shapes\n        self.specialize_int_float = specialize_int_float\n        self.verbose = verbose\n\n    def activate(self) -> None:\n        torchdynamo.config.capture_scalar_outputs = self.capture_scalar_outputs\n        torchdynamo.config.guard_nn_modules = self.guard_nn_modules\n        torchdynamo.config.dynamic_shapes = self.dynamic_shapes\n        torchdynamo.config.specialize_int_float = self.specialize_int_float\n        torchdynamo.config.verbose = self.verbose\n\n    def deactivate(self) -> None:\n        torchdynamo.config.capture_scalar_outputs = True\n        torchdynamo.config.guard_nn_modules = True\n        torchdynamo.config.dynamic_shapes = True\n        torchdynamo.config.specialize_int_float = True\n        torchdynamo.config.verbose = True\n\n\n@contextmanager\ndef using_config(config: DynamoConfig) -> Generator[DynamoConfig, None, None]:\n    config.activate()\n    try:\n        yield config\n    finally:\n        config.deactivate()\n\n\n@contextmanager\ndef setting_python_recursive_limit(limit: int = 10000) -> Generator[None, None, None]:\n    \"\"\"\n    Temporarily increase the python interpreter stack recursion limit.\n    This is mostly used for pickling large scale modules.\n    \"\"\"\n    default = sys.getrecursionlimit()\n    if limit > default:\n        sys.setrecursionlimit(limit)\n    try:\n        yield\n    finally:\n        sys.setrecursionlimit(default)\n\n\n@req_torch_version(\"2.0\")\ndef dynamo_trace(\n    f: Callable[..., Value],\n    # pyre-ignore\n    args: Tuple[Any, ...],\n    aten_graph: bool,\n    tracing_mode: str = \"real\",\n    dynamo_config: Optional[DynamoConfig] = None,\n) -> Tuple[torch.fx.GraphModule, Set]:\n    \"\"\"\n    TODO: Once we fully migrate to torchdynamo frontend, we will remove\n    this config option alltogether.  For now, it helps with quick\n    experiments with playing around with TorchDynamo\n    \"\"\"\n    if dynamo_config is None:\n        dynamo_config = DynamoConfig()\n    with using_config(dynamo_config), setting_python_recursive_limit(2000):\n        torchdynamo.reset()\n        try:\n            return torchdynamo.export(\n                f,\n                *copy.deepcopy(args),\n                aten_graph=aten_graph,\n                tracing_mode=tracing_mode,\n            )\n        except torchdynamo.exc.Unsupported as exc:\n            raise RuntimeError(\n                \"The user code is using a feature we don't support. \"\n                \"Please try torchdynamo.explain() to get possible the reasons\",\n            ) from exc\n        except Exception as exc:\n            raise RuntimeError(\n                \"torchdynamo internal error occured. Please see above stacktrace\"\n            ) from exc\n\n\n@req_torch_version(\"2.0\")\ndef trace(f, args, *rest):\n    graph_module, guards = dynamo_trace(f, args, True, \"symbolic\")\n    return graph_module, guards\n\n\n@req_torch_version(\"2.0\")\ndef opt_trace(f, args, *rest):\n    \"\"\"\n    Optimized trace with necessary passes which re-compose some ops or replace some ops\n    These passes should be general and functional purpose\n    \"\"\"\n    passes_list = [\n        compose_bmm,\n        compose_chunk,\n        compose_getitem_slice,\n        replace_aten_reshape_alias_with_replace,\n        replace_aten_op_with_indices,\n        replace_transpose_mm_op_with_linear,  # after compose_bmm\n        replace_native_layernorm_with_layernorm,\n        remove_ops,\n        replace_builtin_ops,  # after replace_native_layernorm_with_layernorm\n        replace_inplace_ops,  # remove it once functionalization is enabled\n    ]\n\n    fx_module, _ = trace(f, args)\n    print(fx_module.graph)\n    for passes in passes_list:\n        pr: PassResult = passes(fx_module)\n        fx_module = pr.graph_module\n\n    fx_module(*args)\n\n    fx_module = run_const_fold(fx_module)\n    print(fx_module.graph)\n    return fx_module\n", "path": "py/torch_tensorrt/fx/tracer/dispatch_tracer/aten_tracer.py"}, {"content": "from enum import Enum\nfrom typing import List, Callable\nfrom packaging import version\n\n# @manual=//deeplearning/trt/python:py_tensorrt\nimport tensorrt as trt\nimport torch\nfrom functorch import make_fx\nfrom functorch.experimental import functionalize\nfrom torch_tensorrt.fx.passes.lower_basic_pass import (\n    replace_op_with_indices,\n    run_const_fold,\n)\n\nfrom .types import Shape, TRTDataType\n\n\nclass LowerPrecision(Enum):\n    FP32 = \"fp32\"\n    FP16 = \"fp16\"\n    INT8 = \"int8\"\n\n\ndef torch_dtype_to_trt(dtype: torch.dtype) -> TRTDataType:\n    \"\"\"\n    Convert PyTorch data types to TensorRT data types.\n\n    Args:\n        dtype (torch.dtype): A PyTorch data type.\n\n    Returns:\n        The equivalent TensorRT data type.\n    \"\"\"\n    if trt.__version__ >= \"7.0\" and dtype == torch.bool:\n        return trt.bool\n    elif dtype == torch.int8:\n        return trt.int8\n    elif dtype == torch.int32:\n        return trt.int32\n    elif dtype == torch.float16:\n        return trt.float16\n    elif dtype == torch.float32:\n        return trt.float32\n    else:\n        raise TypeError(\"%s is not supported by tensorrt\" % dtype)\n\n\ndef torch_dtype_from_trt(dtype: TRTDataType) -> torch.dtype:\n    \"\"\"\n    Convert TensorRT data types to PyTorch data types.\n\n    Args:\n        dtype (TRTDataType): A TensorRT data type.\n\n    Returns:\n        The equivalent PyTorch data type.\n    \"\"\"\n    if dtype == trt.int8:\n        return torch.int8\n    elif trt.__version__ >= \"7.0\" and dtype == trt.bool:\n        return torch.bool\n    elif dtype == trt.int32:\n        return torch.int32\n    elif dtype == trt.float16:\n        return torch.float16\n    elif dtype == trt.float32:\n        return torch.float32\n    else:\n        raise TypeError(\"%s is not supported by torch\" % dtype)\n\n\ndef get_dynamic_dims(shape: Shape) -> List[int]:\n    \"\"\"\n    This function finds the dynamic dimensions in the given\n    shape. A dimension is dynamic if it's -1.\n\n    Args:\n        shape (Shape): A sequence of integer that represents\n            the shape of a tensor.\n\n    Returns:\n        A list of integers contains all the dynamic dimensions\n        in the given shape\n    \"\"\"\n    dynamic_dims = []\n\n    for i, s in enumerate(shape):\n        if s == -1:\n            dynamic_dims.append(i)\n\n    return dynamic_dims\n\n\ndef proxytensor_trace(mod, inputs):\n\n    mod.eval()\n\n    def f(*inp):\n        return mod(*inp)\n\n    mod = make_fx(functionalize(f))(*inputs)\n\n    # Remove const operation. For ex, nn.Linear has transpose operation on weight\n    mod.graph.eliminate_dead_code()\n    mod = run_const_fold(mod)\n    mod = replace_op_with_indices(mod)\n    return mod\n\n\ndef req_torch_version(min_torch_version: str = \"2.dev\"):\n    \"\"\"\n    Create a decorator which verifies the Torch version installed\n    against a specified version range\n\n    Args:\n        min_torch_version (str): The minimum required Torch version\n        for the decorated function to work properly\n\n    Returns:\n        A decorator which raises a descriptive error message if\n        an unsupported Torch version is used\n    \"\"\"\n\n    def nested_decorator(f: Callable):\n        def function_wrapper(*args, **kwargs):\n            # Parse minimum and current Torch versions\n            min_version = version.parse(min_torch_version)\n            current_version = version.parse(torch.__version__)\n\n            if current_version < min_version:\n                raise AssertionError(\n                    f\"Expected Torch version {min_torch_version} or greater, \"\n                    + f\"when calling {f}. Detected version {torch.__version__}\"\n                )\n            else:\n                return f(*args, **kwargs)\n\n        return function_wrapper\n\n    return nested_decorator\n", "path": "py/torch_tensorrt/fx/utils.py"}]}