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