feat: support aten.atan2.out converter (#2829)

Author: chohk88

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -1533,6 +1533,22 @@ def aten_ops_atan2(
     )
 
 
+@dynamo_tensorrt_converter(torch.ops.aten.atan2.out)
+def aten_ops_atan2_out(
+    ctx: ConversionContext,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> TRTTensor:
+    input, other = args[0], args[1]
+    # out = kwargs.get("out"),
+
+    out_return = impl.elementwise.atan2(ctx, target, SourceIR.ATEN, name, input, other)
+
+    return out_return
+
+
 @dynamo_tensorrt_converter(torch.ops.aten.ceil.default)
 def aten_ops_ceil(
     ctx: ConversionContext,

tests/py/dynamo/conversion/test_atan2_aten.py

@@ -108,7 +108,7 @@ def forward(self, lhs_val, rhs_val):
         ]
     )
     def test_atan2_zero(self, dtype, x_val, y_val):
-        class Atan2(nn.Module):
+        class atan2(nn.Module):
             def forward(self, lhs_val, rhs_val):
                 return torch.ops.aten.atan2.default(lhs_val, rhs_val)
 
@@ -123,7 +123,33 @@ def forward(self, lhs_val, rhs_val):
         ]
 
         self.run_test(
-            Atan2(),
+            atan2(),
+            inputs,
+        )
+
+
+class TestAtan2OutConverter(DispatchTestCase):
+    @parameterized.expand(
+        [
+            ((10,), (5,), torch.float),
+            ((10,), (10,), torch.float),
+        ]
+    )
+    def test_atan2_float(self, input_shape, out_shape, dtype):
+        class atan2_out(nn.Module):
+            def forward(self, lhs_val, rhs_val, out):
+                return torch.ops.aten.atan2.out(lhs_val, rhs_val, out=out)
+
+        out = torch.empty(out_shape)
+
+        inputs = [
+            torch.randn(input_shape, dtype=dtype),
+            torch.randn(input_shape, dtype=dtype),
+            out,
+        ]
+
+        self.run_test(
+            atan2_out(),
             inputs,
         )