We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Add support for `aten.select_scatter`


--- 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/dynamo/lowering/_decompositions.py`
Content:
```
1 import logging
2 from typing import Any, Callable, Dict, List, Optional
3 
4 import torch
5 from torch._decomp import register_decomposition
6 from torch._ops import OpOverload
7 from torch_tensorrt.dynamo.conversion.converter_utils import get_positive_dim
8 
9 from ._decomposition_groups import (
10     ENABLED_TORCH_DECOMPOSITIONS,
11     TORCH_TRT_DECOMPOSITIONS,
12     _core_aten_decompositions,
13     aten,
14     torch_disabled_decompositions,
15     torch_enabled_decompositions,
16 )
17 
18 logger = logging.getLogger(__name__)
19 
20 
21 def register_torch_trt_decomposition(
22     aten_op: OpOverload, registry: Optional[Any] = None
23 ) -> Callable[[Any], Any]:
24     """Checks if the decomposition already exists in one of the sets
25     Registers the decomposition via the Torch utility
26 
27     Alerts the user if the decomposition already exists, before registering
28     Throws an AssertionError if the user attempts to register a decomposition
29     which is present in the set of explicitly disabled decompositions
30     """
31     if aten_op in torch_enabled_decompositions:
32         logger.warning(
33             f"Detected custom decomposition for {aten_op}, which conflicts "
34             "with an existing Torch decomposition in torch_enabled_decompositions. "
35             "The custom implementation will take precedence."
36         )
37     elif aten_op in torch_disabled_decompositions:
38         logger.info(
39             f"Detected custom decomposition for {aten_op}, which is present "
40             "in torch_disabled_decompositions."
41         )
42 
43     # Conflicts with _core_aten_decompositions will only occur if
44     # enable_experimental_decompositions is True in get_decompositions
45     if aten_op in _core_aten_decompositions:
46         logger.debug(
47             f"Detected custom decomposition for {aten_op}, which conflicts "
48             "with an existing Torch decomposition in core_aten_decompositions. "
49             "The custom implementation will take precedence."
50         )
51 
52     def register(fn: Callable[[Any], Any]) -> Any:
53         return register_decomposition(aten_op=aten_op, registry=registry)(fn)
54 
55     return register
56 
57 
58 def replace_inplace_op(aten_op: OpOverload, outplace_op: OpOverload) -> Any:
59     """Replace inplace operation with functional equivalent
60     Adapted from:
61     https://github.com/pytorch/pytorch/blob/3344d79e3f732dadd5c85b99a7aa1a022f187929/torch/_decomp/decompositions.py#L3355-L3361
62     """
63 
64     @register_torch_trt_decomposition(aten_op, registry=TORCH_TRT_DECOMPOSITIONS)
65     def inplace_op(*args, **kwargs):  # type: ignore
66         out = outplace_op(*args, **kwargs)
67         return args[0].copy_(out)
68 
69     return inplace_op
70 
71 
72 replace_inplace_op(aten.add_, aten.add)
73 replace_inplace_op(aten.addbmm_, aten.addbmm)
74 replace_inplace_op(aten.addmm_, aten.addmm)
75 replace_inplace_op(aten.addmv_, aten.addmv)
76 replace_inplace_op(aten.baddbmm_, aten.baddbmm)
77 replace_inplace_op(aten.cumprod_, aten.cumprod)
78 replace_inplace_op(aten.index_put_, aten.index_put)
79 replace_inplace_op(aten.index_reduce_, aten.index_reduce)
80 replace_inplace_op(aten.relu_, aten.relu)
81 replace_inplace_op(aten.round_, aten.round)
82 replace_inplace_op(aten.scatter_, aten.scatter)
83 replace_inplace_op(aten.scatter_add_, aten.scatter_add)
84 replace_inplace_op(aten.scatter_reduce_, aten.scatter_reduce)
85 
86 
87 @register_torch_trt_decomposition(aten.rsqrt, registry=TORCH_TRT_DECOMPOSITIONS)
88 def rsqrt_replacement(*args, **kwargs) -> torch.Tensor:  # type: ignore
89     return torch.reciprocal(torch.sqrt(*args, **kwargs))
90 
91 
92 @register_torch_trt_decomposition(aten._unsafe_view, registry=TORCH_TRT_DECOMPOSITIONS)
93 def unsafe_view_replacement(x: torch.Tensor, *args, **kwargs) -> torch.Tensor:  # type: ignore
94     return torch.reshape(x, *args, **kwargs)
95 
96 
97 @register_torch_trt_decomposition(
98     torch.ops.aten.lift_fresh_copy, registry=TORCH_TRT_DECOMPOSITIONS
99 )
100 def lift_fresh_copy_replacement(x: torch.Tensor) -> torch.Tensor:
101     return x
102 
103 
104 @register_torch_trt_decomposition(aten.alias, registry=TORCH_TRT_DECOMPOSITIONS)
105 def alias_replacement(x: torch.Tensor) -> torch.Tensor:
106     return x
107 
108 
109 @register_torch_trt_decomposition(
110     torch.ops.aten.reciprocal.default, registry=TORCH_TRT_DECOMPOSITIONS
111 )
112 def reciprocal_replacement(
113     input_: torch.Tensor,
114 ) -> torch.Tensor:
115     return torch.div(1, input_)
116 
117 
118 @register_torch_trt_decomposition(
119     torch.ops.prims.var.default, registry=TORCH_TRT_DECOMPOSITIONS
120 )
121 def var_decomposition(
122     input_tensor: torch.Tensor,
123     dims: Optional[List[int]],
124     correction: int,
125     output_dtype: Optional[torch.dtype] = None,
126 ) -> torch.Tensor:
127     if dims is None:
128         dims = []
129 
130     # If the dimensions are empty, variance is taken over all dimensions
131     if isinstance(dims, (tuple, list)) and len(dims) == 0:
132         N = input_tensor.numel()
133     # Otherwise, the number of samples is the product of the dimensions reduced over
134     else:
135         N = 1
136         for dim_i in dims:
137             N *= input_tensor.shape[dim_i]
138 
139     # Compute the mean, difference, and correction term as per the formula:
140     # https://pytorch.org/docs/stable/generated/torch.var.html
141 
142     # Additionally, prims does not support keepdim, and so we only keep dimensions
143     # on the first reduction, then remove it for the second
144     sample_mean = torch.mean(input_tensor, dims, keepdim=True)
145     diff = input_tensor - sample_mean
146     squared_diff = diff * diff
147     variance_unnormalized = torch.sum(squared_diff, dims, keepdim=False)
148 
149     if correction is None:
150         correction_term = float(N - 1)
151     elif isinstance(correction, int):
152         correction_term = float(N - correction)
153     elif isinstance(correction, float):
154         correction_term = float(N) - correction
155     else:
156         raise RuntimeError("correction must be int or float")
157 
158     if correction_term <= 0:
159         raise RuntimeError(f"correction term was non-positive, got: {correction_term}")
160 
161     variance = variance_unnormalized / correction_term
162 
163     return variance
164 
165 
166 @register_torch_trt_decomposition(
167     torch.ops.aten.empty_permuted.default, registry=TORCH_TRT_DECOMPOSITIONS
168 )
169 def empty_permuted_decomposition(*args, **kwargs) -> torch.Tensor:
170     empty_size = args[0]
171     empty_permute = args[1]
172     perm = [0] * len(empty_size)
173     for permute_index, permute_element in enumerate(empty_permute):
174         perm[permute_element] = permute_index
175     return torch.empty([empty_size[l] for l in empty_permute], **kwargs).permute(perm)
176 
177 
178 @register_torch_trt_decomposition(
179     torch.ops.aten.slice_scatter.default, registry=TORCH_TRT_DECOMPOSITIONS
180 )
181 def slice_scatter_decomposition(
182     input_tensor: torch.Tensor,
183     src_tensor: torch.Tensor,
184     dim: int,
185     start: Optional[int] = None,
186     end: Optional[int] = None,
187     step: Optional[int] = None,
188 ):
189     dim_size = input_tensor.shape[dim]
190     start = get_positive_dim(start, input_tensor.shape[dim])
191     if end is None:
192         end = dim_size
193     end = get_positive_dim(end, input_tensor.shape[dim])
194     if step is None:
195         step = 1
196 
197     src_dim = src_tensor.shape
198     # step == 0 is not a valid torch case
199     # also src_dim should be equal to slice dimension
200 
201     if start == 0 and end == dim_size and step == 1:
202         return src_tensor
203 
204     cat_tensors = []
205     index_tensor_shape = []
206     for i, src_each_dim in enumerate(list(src_dim)):
207         if i != dim:
208             index_tensor_shape.append(src_each_dim)
209     for index in range(start, end, step):
210         cat_tensors.append(index * torch.ones(index_tensor_shape, dtype=torch.long))
211     index_tensor = torch.stack(cat_tensors, dim).cuda()
212     output_tensor = torch.scatter(input_tensor, dim, index_tensor, src_tensor)
213     return output_tensor
214 
215 
216 def get_decompositions(
217     enable_experimental_decompositions: bool = False,
218 ) -> Dict[OpOverload, Callable[[Any], Any]]:
219     if enable_experimental_decompositions:
220         CORE_ATEN_DECOMPOSITIONS_FILTERED: Dict[OpOverload, Callable[[Any], Any]] = {
221             decomp: _core_aten_decompositions[decomp]
222             for decomp in _core_aten_decompositions
223             if decomp not in torch_disabled_decompositions
224         }
225         return {**CORE_ATEN_DECOMPOSITIONS_FILTERED, **TORCH_TRT_DECOMPOSITIONS}
226     else:
227         return {**ENABLED_TORCH_DECOMPOSITIONS, **TORCH_TRT_DECOMPOSITIONS}
228
```
--- 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/dynamo/lowering/_decompositions.py b/py/torch_tensorrt/dynamo/lowering/_decompositions.py
--- a/py/torch_tensorrt/dynamo/lowering/_decompositions.py
+++ b/py/torch_tensorrt/dynamo/lowering/_decompositions.py
@@ -213,6 +213,19 @@
     return output_tensor
 
 
+@register_torch_trt_decomposition(
+    torch.ops.aten.select_scatter.default, registry=TORCH_TRT_DECOMPOSITIONS
+)
+def select_scatter_decomposition(
+    input_tensor: torch.Tensor,
+    src_tensor: torch.Tensor,
+    dim: int,
+    index: int,
+) -> torch.Tensor:
+    src_tensor = torch.unsqueeze(src_tensor, dim)
+    return torch.slice_scatter(input_tensor, src_tensor, dim, index, index + 1, 1)
+
+
 def get_decompositions(
     enable_experimental_decompositions: bool = False,
 ) -> Dict[OpOverload, Callable[[Any], Any]]:

{"golden_diff": "diff --git a/py/torch_tensorrt/dynamo/lowering/_decompositions.py b/py/torch_tensorrt/dynamo/lowering/_decompositions.py\n--- a/py/torch_tensorrt/dynamo/lowering/_decompositions.py\n+++ b/py/torch_tensorrt/dynamo/lowering/_decompositions.py\n@@ -213,6 +213,19 @@\n     return output_tensor\n \n \n+@register_torch_trt_decomposition(\n+    torch.ops.aten.select_scatter.default, registry=TORCH_TRT_DECOMPOSITIONS\n+)\n+def select_scatter_decomposition(\n+    input_tensor: torch.Tensor,\n+    src_tensor: torch.Tensor,\n+    dim: int,\n+    index: int,\n+) -> torch.Tensor:\n+    src_tensor = torch.unsqueeze(src_tensor, dim)\n+    return torch.slice_scatter(input_tensor, src_tensor, dim, index, index + 1, 1)\n+\n+\n def get_decompositions(\n     enable_experimental_decompositions: bool = False,\n ) -> Dict[OpOverload, Callable[[Any], Any]]:\n", "issue": "Add support for `aten.select_scatter`\n\n", "before_files": [{"content": "import logging\nfrom typing import Any, Callable, Dict, List, Optional\n\nimport torch\nfrom torch._decomp import register_decomposition\nfrom torch._ops import OpOverload\nfrom torch_tensorrt.dynamo.conversion.converter_utils import get_positive_dim\n\nfrom ._decomposition_groups import (\n    ENABLED_TORCH_DECOMPOSITIONS,\n    TORCH_TRT_DECOMPOSITIONS,\n    _core_aten_decompositions,\n    aten,\n    torch_disabled_decompositions,\n    torch_enabled_decompositions,\n)\n\nlogger = logging.getLogger(__name__)\n\n\ndef register_torch_trt_decomposition(\n    aten_op: OpOverload, registry: Optional[Any] = None\n) -> Callable[[Any], Any]:\n    \"\"\"Checks if the decomposition already exists in one of the sets\n    Registers the decomposition via the Torch utility\n\n    Alerts the user if the decomposition already exists, before registering\n    Throws an AssertionError if the user attempts to register a decomposition\n    which is present in the set of explicitly disabled decompositions\n    \"\"\"\n    if aten_op in torch_enabled_decompositions:\n        logger.warning(\n            f\"Detected custom decomposition for {aten_op}, which conflicts \"\n            \"with an existing Torch decomposition in torch_enabled_decompositions. \"\n            \"The custom implementation will take precedence.\"\n        )\n    elif aten_op in torch_disabled_decompositions:\n        logger.info(\n            f\"Detected custom decomposition for {aten_op}, which is present \"\n            \"in torch_disabled_decompositions.\"\n        )\n\n    # Conflicts with _core_aten_decompositions will only occur if\n    # enable_experimental_decompositions is True in get_decompositions\n    if aten_op in _core_aten_decompositions:\n        logger.debug(\n            f\"Detected custom decomposition for {aten_op}, which conflicts \"\n            \"with an existing Torch decomposition in core_aten_decompositions. \"\n            \"The custom implementation will take precedence.\"\n        )\n\n    def register(fn: Callable[[Any], Any]) -> Any:\n        return register_decomposition(aten_op=aten_op, registry=registry)(fn)\n\n    return register\n\n\ndef replace_inplace_op(aten_op: OpOverload, outplace_op: OpOverload) -> Any:\n    \"\"\"Replace inplace operation with functional equivalent\n    Adapted from:\n    https://github.com/pytorch/pytorch/blob/3344d79e3f732dadd5c85b99a7aa1a022f187929/torch/_decomp/decompositions.py#L3355-L3361\n    \"\"\"\n\n    @register_torch_trt_decomposition(aten_op, registry=TORCH_TRT_DECOMPOSITIONS)\n    def inplace_op(*args, **kwargs):  # type: ignore\n        out = outplace_op(*args, **kwargs)\n        return args[0].copy_(out)\n\n    return inplace_op\n\n\nreplace_inplace_op(aten.add_, aten.add)\nreplace_inplace_op(aten.addbmm_, aten.addbmm)\nreplace_inplace_op(aten.addmm_, aten.addmm)\nreplace_inplace_op(aten.addmv_, aten.addmv)\nreplace_inplace_op(aten.baddbmm_, aten.baddbmm)\nreplace_inplace_op(aten.cumprod_, aten.cumprod)\nreplace_inplace_op(aten.index_put_, aten.index_put)\nreplace_inplace_op(aten.index_reduce_, aten.index_reduce)\nreplace_inplace_op(aten.relu_, aten.relu)\nreplace_inplace_op(aten.round_, aten.round)\nreplace_inplace_op(aten.scatter_, aten.scatter)\nreplace_inplace_op(aten.scatter_add_, aten.scatter_add)\nreplace_inplace_op(aten.scatter_reduce_, aten.scatter_reduce)\n\n\n@register_torch_trt_decomposition(aten.rsqrt, registry=TORCH_TRT_DECOMPOSITIONS)\ndef rsqrt_replacement(*args, **kwargs) -> torch.Tensor:  # type: ignore\n    return torch.reciprocal(torch.sqrt(*args, **kwargs))\n\n\n@register_torch_trt_decomposition(aten._unsafe_view, registry=TORCH_TRT_DECOMPOSITIONS)\ndef unsafe_view_replacement(x: torch.Tensor, *args, **kwargs) -> torch.Tensor:  # type: ignore\n    return torch.reshape(x, *args, **kwargs)\n\n\n@register_torch_trt_decomposition(\n    torch.ops.aten.lift_fresh_copy, registry=TORCH_TRT_DECOMPOSITIONS\n)\ndef lift_fresh_copy_replacement(x: torch.Tensor) -> torch.Tensor:\n    return x\n\n\n@register_torch_trt_decomposition(aten.alias, registry=TORCH_TRT_DECOMPOSITIONS)\ndef alias_replacement(x: torch.Tensor) -> torch.Tensor:\n    return x\n\n\n@register_torch_trt_decomposition(\n    torch.ops.aten.reciprocal.default, registry=TORCH_TRT_DECOMPOSITIONS\n)\ndef reciprocal_replacement(\n    input_: torch.Tensor,\n) -> torch.Tensor:\n    return torch.div(1, input_)\n\n\n@register_torch_trt_decomposition(\n    torch.ops.prims.var.default, registry=TORCH_TRT_DECOMPOSITIONS\n)\ndef var_decomposition(\n    input_tensor: torch.Tensor,\n    dims: Optional[List[int]],\n    correction: int,\n    output_dtype: Optional[torch.dtype] = None,\n) -> torch.Tensor:\n    if dims is None:\n        dims = []\n\n    # If the dimensions are empty, variance is taken over all dimensions\n    if isinstance(dims, (tuple, list)) and len(dims) == 0:\n        N = input_tensor.numel()\n    # Otherwise, the number of samples is the product of the dimensions reduced over\n    else:\n        N = 1\n        for dim_i in dims:\n            N *= input_tensor.shape[dim_i]\n\n    # Compute the mean, difference, and correction term as per the formula:\n    # https://pytorch.org/docs/stable/generated/torch.var.html\n\n    # Additionally, prims does not support keepdim, and so we only keep dimensions\n    # on the first reduction, then remove it for the second\n    sample_mean = torch.mean(input_tensor, dims, keepdim=True)\n    diff = input_tensor - sample_mean\n    squared_diff = diff * diff\n    variance_unnormalized = torch.sum(squared_diff, dims, keepdim=False)\n\n    if correction is None:\n        correction_term = float(N - 1)\n    elif isinstance(correction, int):\n        correction_term = float(N - correction)\n    elif isinstance(correction, float):\n        correction_term = float(N) - correction\n    else:\n        raise RuntimeError(\"correction must be int or float\")\n\n    if correction_term <= 0:\n        raise RuntimeError(f\"correction term was non-positive, got: {correction_term}\")\n\n    variance = variance_unnormalized / correction_term\n\n    return variance\n\n\n@register_torch_trt_decomposition(\n    torch.ops.aten.empty_permuted.default, registry=TORCH_TRT_DECOMPOSITIONS\n)\ndef empty_permuted_decomposition(*args, **kwargs) -> torch.Tensor:\n    empty_size = args[0]\n    empty_permute = args[1]\n    perm = [0] * len(empty_size)\n    for permute_index, permute_element in enumerate(empty_permute):\n        perm[permute_element] = permute_index\n    return torch.empty([empty_size[l] for l in empty_permute], **kwargs).permute(perm)\n\n\n@register_torch_trt_decomposition(\n    torch.ops.aten.slice_scatter.default, registry=TORCH_TRT_DECOMPOSITIONS\n)\ndef slice_scatter_decomposition(\n    input_tensor: torch.Tensor,\n    src_tensor: torch.Tensor,\n    dim: int,\n    start: Optional[int] = None,\n    end: Optional[int] = None,\n    step: Optional[int] = None,\n):\n    dim_size = input_tensor.shape[dim]\n    start = get_positive_dim(start, input_tensor.shape[dim])\n    if end is None:\n        end = dim_size\n    end = get_positive_dim(end, input_tensor.shape[dim])\n    if step is None:\n        step = 1\n\n    src_dim = src_tensor.shape\n    # step == 0 is not a valid torch case\n    # also src_dim should be equal to slice dimension\n\n    if start == 0 and end == dim_size and step == 1:\n        return src_tensor\n\n    cat_tensors = []\n    index_tensor_shape = []\n    for i, src_each_dim in enumerate(list(src_dim)):\n        if i != dim:\n            index_tensor_shape.append(src_each_dim)\n    for index in range(start, end, step):\n        cat_tensors.append(index * torch.ones(index_tensor_shape, dtype=torch.long))\n    index_tensor = torch.stack(cat_tensors, dim).cuda()\n    output_tensor = torch.scatter(input_tensor, dim, index_tensor, src_tensor)\n    return output_tensor\n\n\ndef get_decompositions(\n    enable_experimental_decompositions: bool = False,\n) -> Dict[OpOverload, Callable[[Any], Any]]:\n    if enable_experimental_decompositions:\n        CORE_ATEN_DECOMPOSITIONS_FILTERED: Dict[OpOverload, Callable[[Any], Any]] = {\n            decomp: _core_aten_decompositions[decomp]\n            for decomp in _core_aten_decompositions\n            if decomp not in torch_disabled_decompositions\n        }\n        return {**CORE_ATEN_DECOMPOSITIONS_FILTERED, **TORCH_TRT_DECOMPOSITIONS}\n    else:\n        return {**ENABLED_TORCH_DECOMPOSITIONS, **TORCH_TRT_DECOMPOSITIONS}\n", "path": "py/torch_tensorrt/dynamo/lowering/_decompositions.py"}], "after_files": [{"content": "import logging\nfrom typing import Any, Callable, Dict, List, Optional\n\nimport torch\nfrom torch._decomp import register_decomposition\nfrom torch._ops import OpOverload\nfrom torch_tensorrt.dynamo.conversion.converter_utils import get_positive_dim\n\nfrom ._decomposition_groups import (\n    ENABLED_TORCH_DECOMPOSITIONS,\n    TORCH_TRT_DECOMPOSITIONS,\n    _core_aten_decompositions,\n    aten,\n    torch_disabled_decompositions,\n    torch_enabled_decompositions,\n)\n\nlogger = logging.getLogger(__name__)\n\n\ndef register_torch_trt_decomposition(\n    aten_op: OpOverload, registry: Optional[Any] = None\n) -> Callable[[Any], Any]:\n    \"\"\"Checks if the decomposition already exists in one of the sets\n    Registers the decomposition via the Torch utility\n\n    Alerts the user if the decomposition already exists, before registering\n    Throws an AssertionError if the user attempts to register a decomposition\n    which is present in the set of explicitly disabled decompositions\n    \"\"\"\n    if aten_op in torch_enabled_decompositions:\n        logger.warning(\n            f\"Detected custom decomposition for {aten_op}, which conflicts \"\n            \"with an existing Torch decomposition in torch_enabled_decompositions. \"\n            \"The custom implementation will take precedence.\"\n        )\n    elif aten_op in torch_disabled_decompositions:\n        logger.info(\n            f\"Detected custom decomposition for {aten_op}, which is present \"\n            \"in torch_disabled_decompositions.\"\n        )\n\n    # Conflicts with _core_aten_decompositions will only occur if\n    # enable_experimental_decompositions is True in get_decompositions\n    if aten_op in _core_aten_decompositions:\n        logger.debug(\n            f\"Detected custom decomposition for {aten_op}, which conflicts \"\n            \"with an existing Torch decomposition in core_aten_decompositions. \"\n            \"The custom implementation will take precedence.\"\n        )\n\n    def register(fn: Callable[[Any], Any]) -> Any:\n        return register_decomposition(aten_op=aten_op, registry=registry)(fn)\n\n    return register\n\n\ndef replace_inplace_op(aten_op: OpOverload, outplace_op: OpOverload) -> Any:\n    \"\"\"Replace inplace operation with functional equivalent\n    Adapted from:\n    https://github.com/pytorch/pytorch/blob/3344d79e3f732dadd5c85b99a7aa1a022f187929/torch/_decomp/decompositions.py#L3355-L3361\n    \"\"\"\n\n    @register_torch_trt_decomposition(aten_op, registry=TORCH_TRT_DECOMPOSITIONS)\n    def inplace_op(*args, **kwargs):  # type: ignore\n        out = outplace_op(*args, **kwargs)\n        return args[0].copy_(out)\n\n    return inplace_op\n\n\nreplace_inplace_op(aten.add_, aten.add)\nreplace_inplace_op(aten.addbmm_, aten.addbmm)\nreplace_inplace_op(aten.addmm_, aten.addmm)\nreplace_inplace_op(aten.addmv_, aten.addmv)\nreplace_inplace_op(aten.baddbmm_, aten.baddbmm)\nreplace_inplace_op(aten.cumprod_, aten.cumprod)\nreplace_inplace_op(aten.index_put_, aten.index_put)\nreplace_inplace_op(aten.index_reduce_, aten.index_reduce)\nreplace_inplace_op(aten.relu_, aten.relu)\nreplace_inplace_op(aten.round_, aten.round)\nreplace_inplace_op(aten.scatter_, aten.scatter)\nreplace_inplace_op(aten.scatter_add_, aten.scatter_add)\nreplace_inplace_op(aten.scatter_reduce_, aten.scatter_reduce)\n\n\n@register_torch_trt_decomposition(aten.rsqrt, registry=TORCH_TRT_DECOMPOSITIONS)\ndef rsqrt_replacement(*args, **kwargs) -> torch.Tensor:  # type: ignore\n    return torch.reciprocal(torch.sqrt(*args, **kwargs))\n\n\n@register_torch_trt_decomposition(aten._unsafe_view, registry=TORCH_TRT_DECOMPOSITIONS)\ndef unsafe_view_replacement(x: torch.Tensor, *args, **kwargs) -> torch.Tensor:  # type: ignore\n    return torch.reshape(x, *args, **kwargs)\n\n\n@register_torch_trt_decomposition(\n    torch.ops.aten.lift_fresh_copy, registry=TORCH_TRT_DECOMPOSITIONS\n)\ndef lift_fresh_copy_replacement(x: torch.Tensor) -> torch.Tensor:\n    return x\n\n\n@register_torch_trt_decomposition(aten.alias, registry=TORCH_TRT_DECOMPOSITIONS)\ndef alias_replacement(x: torch.Tensor) -> torch.Tensor:\n    return x\n\n\n@register_torch_trt_decomposition(\n    torch.ops.aten.reciprocal.default, registry=TORCH_TRT_DECOMPOSITIONS\n)\ndef reciprocal_replacement(\n    input_: torch.Tensor,\n) -> torch.Tensor:\n    return torch.div(1, input_)\n\n\n@register_torch_trt_decomposition(\n    torch.ops.prims.var.default, registry=TORCH_TRT_DECOMPOSITIONS\n)\ndef var_decomposition(\n    input_tensor: torch.Tensor,\n    dims: Optional[List[int]],\n    correction: int,\n    output_dtype: Optional[torch.dtype] = None,\n) -> torch.Tensor:\n    if dims is None:\n        dims = []\n\n    # If the dimensions are empty, variance is taken over all dimensions\n    if isinstance(dims, (tuple, list)) and len(dims) == 0:\n        N = input_tensor.numel()\n    # Otherwise, the number of samples is the product of the dimensions reduced over\n    else:\n        N = 1\n        for dim_i in dims:\n            N *= input_tensor.shape[dim_i]\n\n    # Compute the mean, difference, and correction term as per the formula:\n    # https://pytorch.org/docs/stable/generated/torch.var.html\n\n    # Additionally, prims does not support keepdim, and so we only keep dimensions\n    # on the first reduction, then remove it for the second\n    sample_mean = torch.mean(input_tensor, dims, keepdim=True)\n    diff = input_tensor - sample_mean\n    squared_diff = diff * diff\n    variance_unnormalized = torch.sum(squared_diff, dims, keepdim=False)\n\n    if correction is None:\n        correction_term = float(N - 1)\n    elif isinstance(correction, int):\n        correction_term = float(N - correction)\n    elif isinstance(correction, float):\n        correction_term = float(N) - correction\n    else:\n        raise RuntimeError(\"correction must be int or float\")\n\n    if correction_term <= 0:\n        raise RuntimeError(f\"correction term was non-positive, got: {correction_term}\")\n\n    variance = variance_unnormalized / correction_term\n\n    return variance\n\n\n@register_torch_trt_decomposition(\n    torch.ops.aten.empty_permuted.default, registry=TORCH_TRT_DECOMPOSITIONS\n)\ndef empty_permuted_decomposition(*args, **kwargs) -> torch.Tensor:\n    empty_size = args[0]\n    empty_permute = args[1]\n    perm = [0] * len(empty_size)\n    for permute_index, permute_element in enumerate(empty_permute):\n        perm[permute_element] = permute_index\n    return torch.empty([empty_size[l] for l in empty_permute], **kwargs).permute(perm)\n\n\n@register_torch_trt_decomposition(\n    torch.ops.aten.slice_scatter.default, registry=TORCH_TRT_DECOMPOSITIONS\n)\ndef slice_scatter_decomposition(\n    input_tensor: torch.Tensor,\n    src_tensor: torch.Tensor,\n    dim: int,\n    start: Optional[int] = None,\n    end: Optional[int] = None,\n    step: Optional[int] = None,\n):\n    dim_size = input_tensor.shape[dim]\n    start = get_positive_dim(start, input_tensor.shape[dim])\n    if end is None:\n        end = dim_size\n    end = get_positive_dim(end, input_tensor.shape[dim])\n    if step is None:\n        step = 1\n\n    src_dim = src_tensor.shape\n    # step == 0 is not a valid torch case\n    # also src_dim should be equal to slice dimension\n\n    if start == 0 and end == dim_size and step == 1:\n        return src_tensor\n\n    cat_tensors = []\n    index_tensor_shape = []\n    for i, src_each_dim in enumerate(list(src_dim)):\n        if i != dim:\n            index_tensor_shape.append(src_each_dim)\n    for index in range(start, end, step):\n        cat_tensors.append(index * torch.ones(index_tensor_shape, dtype=torch.long))\n    index_tensor = torch.stack(cat_tensors, dim).cuda()\n    output_tensor = torch.scatter(input_tensor, dim, index_tensor, src_tensor)\n    return output_tensor\n\n\n@register_torch_trt_decomposition(\n    torch.ops.aten.select_scatter.default, registry=TORCH_TRT_DECOMPOSITIONS\n)\ndef select_scatter_decomposition(\n    input_tensor: torch.Tensor,\n    src_tensor: torch.Tensor,\n    dim: int,\n    index: int,\n) -> torch.Tensor:\n    src_tensor = torch.unsqueeze(src_tensor, dim)\n    return torch.slice_scatter(input_tensor, src_tensor, dim, index, index + 1, 1)\n\n\ndef get_decompositions(\n    enable_experimental_decompositions: bool = False,\n) -> Dict[OpOverload, Callable[[Any], Any]]:\n    if enable_experimental_decompositions:\n        CORE_ATEN_DECOMPOSITIONS_FILTERED: Dict[OpOverload, Callable[[Any], Any]] = {\n            decomp: _core_aten_decompositions[decomp]\n            for decomp in _core_aten_decompositions\n            if decomp not in torch_disabled_decompositions\n        }\n        return {**CORE_ATEN_DECOMPOSITIONS_FILTERED, **TORCH_TRT_DECOMPOSITIONS}\n    else:\n        return {**ENABLED_TORCH_DECOMPOSITIONS, **TORCH_TRT_DECOMPOSITIONS}\n", "path": "py/torch_tensorrt/dynamo/lowering/_decompositions.py"}]}