test cases for boolean input

Author: bowang007

py/torch_tensorrt/dynamo/conversion/impl/reduce.py

@@ -220,12 +220,13 @@ def any(
     dim: Sequence[int] = [],
     keepdim: bool = False,
 ) -> TRTTensor:
+    num_out = cast_trt_tensor(ctx, input_val, trt.float32, f"{name}_cast")
     abs_out = impl.unary.abs(
         ctx,
         target,
         source_ir,
         f"{name}_abs",
-        input_val,
+        num_out,
     )
     max_out = amax(ctx, target, source_ir, f"{name}_amax", abs_out, dim, keepdim)
 

tests/py/dynamo/conversion/test_any.py

@@ -126,6 +126,69 @@ def forward(self, x):
             output_dtypes=[torch.bool],
         )
 
+    @parameterized.expand(
+        [
+            ((2, 3, 4), torch.int, -5, 0),
+            ((6, 7, 5, 4, 5), torch.int, -5, 5),
+            ((1, 5, 2, 1), torch.int, -5, 5),
+        ]
+    )
+    def test_any_default_bool_dtype(self, input_shape, dtype, low, high):
+        class Any(nn.Module):
+            def forward(self, x):
+                return torch.ops.aten.any.default(x)
+
+        inputs = [torch.randint(low, high, input_shape, dtype=dtype).bool()]
+        self.run_test(
+            Any(),
+            inputs,
+            output_dtypes=[torch.bool],
+        )
+
+    @parameterized.expand(
+        [
+            ((3, 2, 4), 1, True, torch.int, 0, 5),
+            ((2, 3, 4, 5), 3, True, torch.int, -10, 10),
+            ((2, 3, 4, 5), 2, False, torch.int32, -5, 0),
+            ((6, 7, 5, 4, 5), 4, False, torch.int32, -5, 5),
+            ((1, 5, 2, 1), -4, False, torch.int32, -5, 5),
+        ]
+    )
+    def test_any_dim_bool_dtype(self, input_shape, dim, keep_dims, dtype, low, high):
+        class AnyDim(nn.Module):
+            def forward(self, x):
+                return torch.ops.aten.any.dim(x, dim, keep_dims)
+
+        inputs = [torch.randint(low, high, input_shape, dtype=dtype).bool()]
+        self.run_test(
+            AnyDim(),
+            inputs,
+            output_dtypes=[torch.bool],
+        )
+
+    @parameterized.expand(
+        [
+            ((3, 2, 4), [1], True, torch.int, 0, 5),
+            ((2, 1, 4, 5), [0, 3], True, torch.int, -10, 10),
+            ((2, 3, 4, 5), [0, 1, 2, 3], False, torch.int32, -5, 0),
+            ((6, 7, 5, 4, 5), [1, 3, 4], False, torch.int32, -5, 5),
+            ((1, 5, 2, 1), [-3, -1], False, torch.int32, -5, 5),
+        ]
+    )
+    def test_any_dims_tuple_bool_dtype(
+        self, input_shape, dims, keep_dims, dtype, low, high
+    ):
+        class AnyDims(nn.Module):
+            def forward(self, x):
+                return torch.ops.aten.any.dims(x, dims, keep_dims)
+
+        inputs = [torch.randint(low, high, input_shape, dtype=dtype).bool()]
+        self.run_test(
+            AnyDims(),
+            inputs,
+            output_dtypes=[torch.bool],
+        )
+
 
 if __name__ == "__main__":
     run_tests()