add atanh on diopi torch impl

DeepLink-org · Aug 14, 2024 · 89168cd · 89168cd
1 parent 6f4e9d6
commit 89168cd
Show file tree

Hide file tree

Showing 4 changed files with 73 additions and 0 deletions.
diff --git a/diopi_test/python/configs/diopi_configs.py b/diopi_test/python/configs/diopi_configs.py
@@ -1383,6 +1383,44 @@
         ),
     ),
 
+    'atanh': dict(
+        name=['atanh'],
+        interface=['torch'],
+        is_inplace=True,
+        saved_args=dict(output=0),
+        dtype=[np.float16, np.float32, np.float64],
+        tensor_para=dict(
+            gen_fn='Genfunc.randn',
+            args=[
+                {
+                    "ins": ['input'],
+                    "shape": ((), (1, ), (1024,), (364800, 4), (2, 128, 3072),
+                              (256, 128, 3, 3),
+                              (2, 31, 512, 6, 40),
+                              (0,), (16, 0), (1, 0, 6)),
+                },
+            ],
+        ),
+    ),
+
+    'atanh_not_float': dict(
+        name=['atanh'],
+        interface=['torch'],
+        dtype=[np.int16, np.int32, np.int64, np.uint8, np.int8, np.bool_],
+        tensor_para=dict(
+            gen_fn='Genfunc.randn',
+            args=[
+                {
+                    "ins": ['input'],
+                    "shape": ((), (1, ), (1024,), (364800, 4), (2, 128, 3072),
+                              (256, 128, 3, 3),
+                              (2, 31, 512, 6, 40),
+                              (0,), (16, 0), (1, 0, 6)),
+                },
+            ],
+        ),
+    ),
+
     'sign': dict(
         name=['sign'],
         interface=['torch'],

diff --git a/diopi_test/python/conformance/diopi_functions.py b/diopi_test/python/conformance/diopi_functions.py
@@ -436,6 +436,9 @@ def asinh(input, inplace=False) -> Tensor:
 def acosh(input, inplace=False) -> Tensor:
     return unary_op(input, inplace, "diopiAcosh", promote_type(input, Dtype.float32))
 
+def atanh(input, inplace=False) -> Tensor:
+    return unary_op(input, inplace, "diopiAtanh", promote_type(input, Dtype.float32))
+
 
 def exp(input, inplace=False) -> Tensor:
     return unary_op(input, inplace, "diopiExp", promote_type(input, Dtype.float32))

diff --git a/impl/torch/functions/functions.cpp b/impl/torch/functions/functions.cpp
@@ -1060,6 +1060,23 @@ diopiError_t diopiAcoshInp(diopiContextHandle_t ctx, diopiTensorHandle_t input)
     return diopiSuccess;
 }
 
+diopiError_t diopiAtanh(diopiContextHandle_t ctx, diopiTensorHandle_t out, diopiConstTensorHandle_t input) {
+    impl::aten::setCurStream(ctx);
+    auto atInput = impl::aten::buildATen(input);
+    auto atOut = impl::aten::buildATen(out);
+    CALL_ATEN_CUDA_FUNC(atanh_out, atOut, atInput);
+
+    return diopiSuccess;
+}
+
+diopiError_t diopiAtanhInp(diopiContextHandle_t ctx, diopiTensorHandle_t input) {
+    impl::aten::setCurStream(ctx);
+    auto atInput = impl::aten::buildATen(input);
+    CALL_ATEN_CUDA_FUNC(atanh_, atInput);
+
+    return diopiSuccess;
+}
+
 diopiError_t diopiSigmoid(diopiContextHandle_t ctx, diopiTensorHandle_t out, diopiConstTensorHandle_t input) {
     impl::aten::setCurStream(ctx);
     auto atInput = impl::aten::buildATen(input);

diff --git a/proto/include/diopi/functions.h b/proto/include/diopi/functions.h
@@ -1007,6 +1007,21 @@ DIOPI_API diopiError_t diopiAcoshInp(diopiContextHandle_t ctx, diopiTensorHandle
  */
 DIOPI_API diopiError_t diopiAcosh(diopiContextHandle_t ctx, diopiTensorHandle_t out, diopiConstTensorHandle_t input);
 
+/**
+ * @brief The in-place version of diopiAtanh().
+ * @param[in] ctx Context environment.
+ * @param[inout] input the input tensor and will be stroed reuslt tensor. type = [float16, float32, float64].
+ */
+DIOPI_API diopiError_t diopiAtanhInp(diopiContextHandle_t ctx, diopiTensorHandle_t input);
+
+/**
+ * @brief Returns a new tensor with the arc hyperbolic tangent of the elements of input.
+ * @param[in] ctx Context environment.
+ * @param[in] input the input tensor. type = [float16, float32, float64].
+ * @param[out] out the output tensor. type = [float16, float32, float64].
+ */
+DIOPI_API diopiError_t diopiAtanh(diopiContextHandle_t ctx, diopiTensorHandle_t out, diopiConstTensorHandle_t input);
+
 /**
  * @brief The in-place version of diopiSigmoid().
  * @param[in] ctx Context environment.