Fixed a typo in the converter. Covered the discontinuous tests

Author: cehongwang

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

@@ -713,7 +713,7 @@ def index_put_converter(
     for dim in range(rank):
         unique_suffix = f"{dim}_{i_idx if dim in I else f_idx}"
         if dim in I:
-            idx_tensor = I_combined[i]
+            idx_tensor = I_combined[i_idx]
             ii_list.append(idx_tensor)
             i_idx += 1
         else:
@@ -771,7 +771,12 @@ def index_put_converter(
         )
     else:  # Non-scalar case
         values_shape = list(values.shape)
-        if K > 0 and N in values_shape:
+        if (
+            K > 0
+            and N in values_shape
+            and (len(F) > 1 and max(F) - min(F) + 1 == len(F))
+        ):
+            # Continuous case
             n_idx = values_shape.index(N)
             permute_order = [n_idx] + [
                 i for i in range(len(values_shape)) if i != n_idx
@@ -817,6 +822,7 @@ def index_put_converter(
                 tuple(broadcast_shape),
             )
         else:
+            # Discontinuous case
             values_shape_padded = [1] * (
                 len(expected_shape) - len(values.shape)
             ) + list(values.shape)
@@ -828,20 +834,20 @@ def index_put_converter(
                     raise ValueError(
                         f"Cannot broadcast {values.shape} to {expected_shape}"
                     )
-            values_reshaped = impl.shuffle.reshape(
-                ctx,
-                target,
-                source_ir,
-                f"{name}_reshape_values",
-                values,
-                tuple(broadcast_shape),
-            )
+            # values_reshaped = impl.shuffle.reshape(
+            #     ctx,
+            #     target,
+            #     source_ir,
+            #     f"{name}_reshape_values",
+            #     values,
+            #     tuple(broadcast_shape),
+            # )
             values_expanded = impl.slice.expand(
                 ctx,
                 target,
                 source_ir,
                 f"{name}_expand_values",
-                values_reshaped,
+                values,
                 expected_shape,
             )
 

tests/py/dynamo/conversion/test_index_put_aten.py

@@ -195,11 +195,43 @@ class TestIndexPutConverter(DispatchTestCase):
                     dtype=torch.int32,
                 ),
             ),
+            # param(
+            #     test_name="4d_indices_none_none_multiple_idx_broadcast_error",
+            #     source_tensor=torch.zeros([1, 2, 5, 3], dtype=torch.float32),
+            #     indices_tensor=(None, None, torch.tensor([0, 1, 2], dtype=torch.int64)),
+            #     value_tensor=torch.randn([2, 3, 3], dtype=torch.float32),
+            # ),
+            param(
+                test_name="discontinuous_test",
+                source_tensor=torch.zeros([2, 4, 4], dtype=torch.float32),
+                indices_tensor=(
+                    torch.tensor([0, 0, 1], dtype=torch.int64),
+                    None,
+                    torch.tensor([0, 0, 1], dtype=torch.int64),
+                ),
+                value_tensor=torch.tensor([2, 3, 3, 4], dtype=torch.float32),
+            ),
             param(
-                test_name="4d_indices_none_none_multiple_idx_broadcast_error",
-                source_tensor=torch.zeros([1, 2, 5, 3], dtype=torch.float32),
-                indices_tensor=(None, None, torch.tensor([0, 1, 2], dtype=torch.int64)),
-                value_tensor=torch.randn([2, 3, 3], dtype=torch.float32),
+                test_name="discontinuous_test_two",
+                source_tensor=torch.zeros([2, 4, 4, 2], dtype=torch.float32),
+                indices_tensor=(
+                    None,
+                    torch.tensor([0, 0, 1, 1], dtype=torch.int64),
+                    None,
+                    torch.tensor([0, 0, 1, 1], dtype=torch.int64),
+                ),
+                value_tensor=torch.tensor([2, 3, 3, 4], dtype=torch.float32),
+            ),
+            param(
+                test_name="continuous_test",
+                source_tensor=torch.zeros([2, 4, 4, 2], dtype=torch.float32),
+                indices_tensor=(
+                    None,
+                    None,
+                    torch.tensor([0, 0, 1, 1], dtype=torch.int64),
+                    torch.tensor([0, 0, 1, 1], dtype=torch.int64),
+                ),
+                value_tensor=torch.tensor([2, 3, 3, 4], dtype=torch.float32),
             ),
             # param(
             #     test_name="2d_indices_accumulate_True",