select_scatter changes

Author: apbose

py/torch_tensorrt/dynamo/lowering/_decompositions.py

@@ -163,7 +163,7 @@ def var_decomposition(
 
 
 @register_torch_trt_decomposition(
-    torch.ops.select_scatter, registry=TORCH_TRT_DECOMPOSITIONS
+    torch.ops.aten.select_scatter.default, registry=TORCH_TRT_DECOMPOSITIONS
 )
 def select_scatter_decomposition(
     input_tensor: torch.Tensor,
@@ -177,21 +177,18 @@ def select_scatter_decomposition(
         raise AssertionError("The index should not be greater than dim")
 
     # expanding the src_tensor to have the same dimension as input_tensor
-    src_tensor = torch.expand(torch.unsqueeze(src_tensor, dim), input_tensor.shape)
     # check if the dimension of the src tensor is same as slice tensor
     select_tensor = torch.select(input_tensor, dim, index)
+
     if select_tensor.shape != src_tensor.shape:
         raise AssertionError(
             "The slice tensor shape should be equal to the src tensor shape"
         )
 
-    # make the index tensor
-    # input_tensor_shape = input_tensor.shape
-    # return torch.where(torch.eq((input_tensor_shape[dim]), index), src_tensor, input_tensor)
-
     unbind_tensors = torch.unbind(input_tensor, dim)
-    unbind_tensors[index] = src_tensor
-    return torch.cat(unbind_tensors, dim)
+    unbind_tensors_list = list(unbind_tensors)
+    unbind_tensors_list[index] = src_tensor
+    return torch.stack(tuple(unbind_tensors_list), dim)
 
 
 def get_decompositions(

tests/py/dynamo/lowering/test_decompositions.py

@@ -1,9 +1,9 @@
 import torch
-import torch_tensorrt
 from torch.testing._internal.common_utils import TestCase, run_tests
 
-from ..testing_utilities import DECIMALS_OF_AGREEMENT, lower_graph_testing
+import torch_tensorrt
 
+from ..testing_utilities import DECIMALS_OF_AGREEMENT, lower_graph_testing
 
 class TestLowering(TestCase):
     def test_lowering_inplace_op(self):
@@ -426,19 +426,18 @@ def __init__(self, *args, **kwargs) -> None:
                 super().__init__(*args, **kwargs)
 
             def forward(self, x, src, dim, index):
-                y = self.select_scatter(x, src, dim, index)
+                y = torch.ops.aten.select_scatter.default(x, src, dim, index)
                 return y
 
         # Operations expected to be removed in the traced graph after decompositions
         expected_ops = {
-            torch.ops.aten.lt.default,
-            torch.ops.aten.expand.default,
-            torch.ops.aten.unsqueeze.default,
-            torch.ops.aten.where.default,
+            torch.ops.aten.slice.Tensor,
+            torch.ops.aten.squeeze.dim,
+            torch.ops.aten.cat.default,
         }
-        unexpected_ops = {torch.ops.aten.select_scatter}
+        unexpected_ops = {torch.ops.aten.select_scatter.default}
 
-        inputs = [torch.randn(2, 2), torch.ones(2)]
+        inputs = [torch.zeros(2, 2).cuda(), torch.ones(2).cuda(), 0, 0]
 
         fx_graph = torch.fx.symbolic_trace(selectScatter())
         unexpected_ops_seen, expected_ops_unseen = lower_graph_testing(