Adding the broadcast check and adding corresponding tests

Author: apbose

py/torch_tensorrt/dynamo/conversion/impl/select.py

@@ -5,6 +5,7 @@
 import tensorrt as trt
 from torch.fx.node import Target
 from torch_tensorrt.dynamo._SourceIR import SourceIR
+from torch_tensorrt.dynamo.conversion.converter_utils import broadcastable
 from torch_tensorrt.dynamo.conversion.impl.elementwise import convert_binary_elementwise
 from torch_tensorrt.dynamo.conversion.impl.shape import get_shape_with_dynamic_shape
 from torch_tensorrt.fx.converters.converter_utils import (
@@ -84,12 +85,21 @@ def index(
     # check if the input is dynamic
     dynamic_shape = has_dynamic_shape(input.shape)
 
+    # here we need to check if all the index are broadcastable
+    # if no, then we need to broadcast
+
+    last_index = None
+    broadcast_shape_len = 0
     for i, ind in enumerate(index):
         if ind is not None:
             _LOGGER.debug(f"Shape of {i} index is {ind.shape}")
             adv_indx_indices.append(i)
             # torch.nn.parameter.Parameter=> torch.Tensor
             ind = get_trt_tensor(network, ind, f"parameter_to_fp32_tensor_{i}")
+            if last_index is not None:
+                if not (broadcastable(ind, last_index)):
+                    assert "The indices should be broadcastable"
+            last_index = ind
             tensor_indices.append(ind)
 
     if not tensor_indices:

tests/py/dynamo/conversion/test_index_aten.py

@@ -86,6 +86,54 @@ def forward(self, x):
             expected_ops={torch.ops.aten.index.Tensor, operator.getitem},
         )
 
+    def test_index_zero_index_one_SD(self):
+        class TestModule(nn.Module):
+            def forward(self, x):
+                index0 = torch.tensor([0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7])
+                index1 = torch.tensor([0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7])
+                indices = [None, index0, index1, None]
+                out = torch.ops.aten.index.Tensor(x, indices)
+                return out
+
+        input = [torch.randn(2, 1280, 8, 8)]
+        self.run_test(
+            TestModule(),
+            input,
+            expected_ops={torch.ops.aten.index.Tensor, operator.getitem},
+        )
+
+    def test_index_zero_index_one_SD_unsqueeze(self):
+        class TestModule(nn.Module):
+            def forward(self, x):
+                index0 = torch.tensor([0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7])
+                index1 = index0.unsqueeze(0).T.long()
+                indices = [None, None, index1, index1]
+                out = torch.ops.aten.index.Tensor(x, indices)
+                return out
+
+        input = [torch.randn(2, 1280, 8, 8)]
+        self.run_test(
+            TestModule(),
+            input,
+            expected_ops={torch.ops.aten.index.Tensor},
+        )
+
+    def test_index_zero_index_one_index_two_SD_unsqueeze(self):
+        class TestModule(nn.Module):
+            def forward(self, x):
+                index0 = torch.tensor([0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7])
+                index1 = index0.unsqueeze(0).T.long()
+                indices = [None, None, index0, index1]
+                out = torch.ops.aten.index.Tensor(x, indices)
+                return out
+
+        input = [torch.randn(2, 1280, 8, 8)]
+        self.run_test(
+            TestModule(),
+            input,
+            expected_ops={torch.ops.aten.index.Tensor},
+        )
+
 
 if __name__ == "__main__":
     run_tests()