feat: support aten.any related converters in dynamo

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -2605,6 +2605,27 @@ def aten_ops_remainder(
     )
 
 
+@dynamo_tensorrt_converter(torch.ops.aten.any.default)
+@dynamo_tensorrt_converter(torch.ops.aten.any.dim)
+@dynamo_tensorrt_converter(torch.ops.aten.any.dims)
+def aten_ops_any(
+    ctx: ConversionContext,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.reduce.any(
+        ctx,
+        target,
+        SourceIR.ATEN,
+        name,
+        args[0],
+        args_bounds_check(args, 1, replacement=None),
+        args_bounds_check(args, 2, replacement=False),
+    )
+
+
 @dynamo_tensorrt_converter(torch.ops.aten._pdist_forward.default)
 @enforce_tensor_types(
     {

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

@@ -3,6 +3,7 @@
 import tensorrt as trt
 from torch.fx.node import Target
 from torch_tensorrt.dynamo._SourceIR import SourceIR
+from torch_tensorrt.dynamo.conversion import impl
 from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext
 from torch_tensorrt.dynamo.conversion.converter_utils import (
     cast_trt_tensor,
@@ -208,3 +209,32 @@ def mean(
     )
     set_layer_name(layer, target, name, source_ir)
     return layer.get_output(0)
+
+
+def any(
+    ctx: ConversionContext,
+    target: Target,
+    source_ir: Optional[SourceIR],
+    name: str,
+    input_val: TRTTensor,
+    dim: Union[int, Optional[Sequence[int]]] = None,
+    keepdim: bool = False,
+) -> TRTTensor:
+    if (isinstance(input_val, TRTTensor)) and (input_val.dtype == trt.bool):
+        input_val = cast_trt_tensor(ctx, input_val, trt.int32, f"{name}_cast")
+
+    abs_out = impl.unary.abs(
+        ctx,
+        target,
+        source_ir,
+        f"{name}_abs",
+        input_val,
+    )
+    if dim is None:
+        dim = []
+    elif isinstance(dim, int):
+        dim = [dim]
+
+    max_out = amax(ctx, target, source_ir, f"{name}_amax", abs_out, dim, keepdim)
+
+    return cast_trt_tensor(ctx, max_out, trt.bool, f"{name}_cast_to_bool")

tests/py/dynamo/conversion/test_any.py

@@ -0,0 +1,194 @@
+import torch
+import torch.nn as nn
+from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests
+
+from .harness import DispatchTestCase
+
+
+class TestAnyConverter(DispatchTestCase):
+    @parameterized.expand(
+        [
+            (
+                "3d",
+                (3, 2, 4),
+            ),
+            (
+                "4d",
+                (2, 3, 4, 5),
+            ),
+            ("5d", (6, 7, 5, 4, 5)),
+        ]
+    )
+    def test_any_default_float_dtype(self, _, input_shape):
+        class Any(nn.Module):
+            def forward(self, x):
+                return torch.ops.aten.any.default(x)
+
+        inputs = [torch.randn(*input_shape)]
+        self.run_test(Any(), inputs, output_dtypes=[torch.bool])
+
+    @parameterized.expand(
+        [
+            ((3, 2, 4), 1, True),
+            ((2, 3, 4, 5), 3, True),
+            ((2, 3, 4, 5), 2, False),
+            ((6, 7, 5, 4, 5), 4, False),
+            ((1, 5, 2, 1), -1, True),
+        ]
+    )
+    def test_any_dim_float_dtype(self, input_shape, dim, keep_dims):
+        class AnyDim(nn.Module):
+            def forward(self, x):
+                return torch.ops.aten.any.dim(x, dim, keep_dims)
+
+        inputs = [torch.randn(*input_shape)]
+        self.run_test(AnyDim(), inputs, output_dtypes=[torch.bool])
+
+    @parameterized.expand(
+        [
+            ((3, 2, 4), [1], True),
+            ((2, 1, 4, 5), [0, 3], True),
+            ((2, 3, 4, 5), [0, 1, 2, 3], False),
+            ((6, 7, 5, 4, 5), [1, 3, 4], False),
+        ]
+    )
+    def test_any_dims_tuple_float_dtype(self, input_shape, dims, keep_dims):
+        class AnyDims(nn.Module):
+            def forward(self, x):
+                return torch.ops.aten.any.dims(x, dims, keep_dims)
+
+        inputs = [torch.randn(*input_shape)]
+        self.run_test(AnyDims(), inputs, output_dtypes=[torch.bool])
+
+    @parameterized.expand(
+        [
+            ((3, 2, 4), torch.int, 0, 5),
+            ((2, 3, 4, 5), torch.int, -10, 10),
+            ((2, 3, 4, 5), torch.int32, -5, 0),
+            ((6, 7, 5, 4, 5), torch.int32, -5, 5),
+            ((1, 5, 2, 1), torch.int32, -5, 5),
+        ]
+    )
+    def test_any_default_int_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)]
+        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_int_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)]
+        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_int_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)]
+        self.run_test(
+            AnyDims(),
+            inputs,
+            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()