From c2dc3bc7994719b4c46f61015db70cb9713455ef Mon Sep 17 00:00:00 2001
From: Hoonkyung Cho <chohk88@gmail.com>
Date: Sat, 11 May 2024 17:14:48 +0900
Subject: [PATCH] feat: support aten.atan2.out converter

---
 .../dynamo/conversion/aten_ops_converters.py  | 16 +++++++++
 .../dynamo/conversion/test_atan2_out_aten.py  | 36 +++++++++++++++++++
 2 files changed, 52 insertions(+)
 create mode 100644 tests/py/dynamo/conversion/test_atan2_out_aten.py

diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
index 1705dd06db..48d61f59ac 100644
--- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
+++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
@@ -1427,6 +1427,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_out_aten.py b/tests/py/dynamo/conversion/test_atan2_out_aten.py
new file mode 100644
index 0000000000..eb3158bb72
--- /dev/null
+++ b/tests/py/dynamo/conversion/test_atan2_out_aten.py
@@ -0,0 +1,36 @@
+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 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,
+        )
+
+
+if __name__ == "__main__":
+    run_tests()