From 4f3b7d9de9c8b9bbce5ce0d08f630d66593767ee Mon Sep 17 00:00:00 2001
From: Hoonkyung Cho <chohk88@gmail.com>
Date: Fri, 24 May 2024 23:55:16 +0900
Subject: [PATCH] feat: support aten.atan2.out converter (#2829)

---
 .../dynamo/conversion/aten_ops_converters.py  | 16 ++++++++++
 tests/py/dynamo/conversion/test_atan2_aten.py | 30 +++++++++++++++++--
 2 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
index 96b86f6640..2937c33cda 100644
--- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
+++ b/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,
diff --git a/tests/py/dynamo/conversion/test_atan2_aten.py b/tests/py/dynamo/conversion/test_atan2_aten.py
index 550ade2970..7b367fcd13 100644
--- a/tests/py/dynamo/conversion/test_atan2_aten.py
+++ b/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,
         )