Supported bool mask indicies

Author: cehongwang

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

@@ -550,6 +550,9 @@ def index_put_converter(
     accumulate: bool = False,
 ) -> TRTTensor:
     # Convert 'input_indices' to TRT tensors (or keep None as is)
+    input_indices = expand_boolean_indices(
+        ctx, target, source_ir, name, input_tensor, input_indices
+    )
     indices: List[Optional[Union[TRTTensor, None]]] = []
     for i, idx in enumerate(input_indices):
         if idx is None:
@@ -828,20 +831,15 @@ def index_put_converter(
             ) + list(values.shape)
             broadcast_shape = []
             for exp_dim, val_dim in zip(expected_shape, values_shape_padded):
-                if val_dim == 1 or exp_dim == val_dim:
+                if val_dim == DYNAMIC_DIM or exp_dim == DYNAMIC_DIM:
+                    broadcast_shape.append(-1)
+                elif val_dim == 1 or exp_dim == val_dim:
                     broadcast_shape.append(exp_dim)
                 else:
                     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_expanded = impl.slice.expand(
                 ctx,
                 target,

tests/py/dynamo/conversion/test_index_put_aten.py

@@ -328,6 +328,43 @@ def forward(self, x, y, z, a, b):
         result = trt_mod(*inputs)
         assert torch.allclose(result, torch_output, atol=1e-4, rtol=1e-4)
 
+    def test_bool_mask_test(self):
+
+        source_tensor = torch.ones([5, 10], dtype=torch.float32).cuda()
+        indices_tensor = torch.tensor([False, False, True, False, True])
+        value_tensor = torch.zeros([2, 10], dtype=torch.float32).cuda()
+
+        dim1 = torch.export.Dim("dim1", min=1, max=5)
+        dim2 = torch.export.Dim("dim2", min=1, max=5)
+
+        class TestIndexPut(torch.nn.Module):
+            def forward(self, source_tensor, indices_tensor, value_tensor):
+                source_tensor[indices_tensor] = value_tensor
+                return source_tensor
+
+        model = TestIndexPut()
+        torch_output = model.forward(source_tensor, indices_tensor, value_tensor)
+
+        ep = torch.export.export(
+            model,
+            (source_tensor, indices_tensor, value_tensor),
+            dynamic_shapes=({0: dim1}, {0: dim1}, {0: dim2}),
+        )
+        with torchtrt.dynamo.Debugger(log_level="debug"):
+            trt_engine = torchtrt.dynamo.compile(
+                ep,
+                inputs=(source_tensor, indices_tensor, value_tensor),
+                enabled_precisions={torch.float32},
+                min_block_size=1,
+                use_explicit_typing=False,
+                use_fp32_acc=False,
+                disable_tf32=True,
+                use_python_runtime=True,
+            )
+        result = trt_engine(source_tensor, indices_tensor, value_tensor)
+
+        torch.allclose(result, torch_output, atol=1e-4, rtol=1e-4)
+
 
 if __name__ == "__main__":
     run_tests()