ENH: add atanh ufunc

.github/workflows/array_api.yml

@@ -49,4 +49,4 @@ jobs:
           # for hypothesis-driven test case generation
           pytest $GITHUB_WORKSPACE/pre_compile_tools/pre_compile_ufuncs.py -s
           # only run a subset of the conformance tests to get started
-          pytest array_api_tests/meta/test_broadcasting.py array_api_tests/meta/test_equality_mapping.py array_api_tests/meta/test_signatures.py array_api_tests/meta/test_special_cases.py array_api_tests/test_constants.py array_api_tests/meta/test_utils.py array_api_tests/test_creation_functions.py::test_ones array_api_tests/test_creation_functions.py::test_ones_like array_api_tests/test_data_type_functions.py::test_result_type array_api_tests/test_operators_and_elementwise_functions.py::test_log10 array_api_tests/test_operators_and_elementwise_functions.py::test_sqrt array_api_tests/test_operators_and_elementwise_functions.py::test_isfinite array_api_tests/test_operators_and_elementwise_functions.py::test_log2 array_api_tests/test_operators_and_elementwise_functions.py::test_log1p array_api_tests/test_operators_and_elementwise_functions.py::test_isinf array_api_tests/test_operators_and_elementwise_functions.py::test_log array_api_tests/test_array_object.py::test_scalar_casting array_api_tests/test_operators_and_elementwise_functions.py::test_sign array_api_tests/test_operators_and_elementwise_functions.py::test_square array_api_tests/test_operators_and_elementwise_functions.py::test_cos array_api_tests/test_operators_and_elementwise_functions.py::test_round array_api_tests/test_operators_and_elementwise_functions.py::test_trunc array_api_tests/test_operators_and_elementwise_functions.py::test_ceil array_api_tests/test_operators_and_elementwise_functions.py::test_floor
+          pytest array_api_tests/meta/test_broadcasting.py array_api_tests/meta/test_equality_mapping.py array_api_tests/meta/test_signatures.py array_api_tests/meta/test_special_cases.py array_api_tests/test_constants.py array_api_tests/meta/test_utils.py array_api_tests/test_creation_functions.py::test_ones array_api_tests/test_creation_functions.py::test_ones_like array_api_tests/test_data_type_functions.py::test_result_type array_api_tests/test_operators_and_elementwise_functions.py::test_log10 array_api_tests/test_operators_and_elementwise_functions.py::test_sqrt array_api_tests/test_operators_and_elementwise_functions.py::test_isfinite array_api_tests/test_operators_and_elementwise_functions.py::test_log2 array_api_tests/test_operators_and_elementwise_functions.py::test_log1p array_api_tests/test_operators_and_elementwise_functions.py::test_isinf array_api_tests/test_operators_and_elementwise_functions.py::test_log array_api_tests/test_array_object.py::test_scalar_casting array_api_tests/test_operators_and_elementwise_functions.py::test_sign array_api_tests/test_operators_and_elementwise_functions.py::test_square array_api_tests/test_operators_and_elementwise_functions.py::test_cos array_api_tests/test_operators_and_elementwise_functions.py::test_round array_api_tests/test_operators_and_elementwise_functions.py::test_trunc array_api_tests/test_operators_and_elementwise_functions.py::test_ceil array_api_tests/test_operators_and_elementwise_functions.py::test_floor array_api_tests/test_operators_and_elementwise_functions.py::test_atanh

pykokkos/__init__.py

@@ -38,6 +38,7 @@
                                  sin,
                                  cos,
                                  tan,
+                                 atanh,
                                  logical_and,
                                  logical_or,
                                  logical_xor,

pykokkos/lib/ufunc_workunits.py

@@ -1,6 +1,28 @@
 import pykokkos as pk
 
 
+@pk.workunit
+def atanh_impl_1d_double(tid: int, view: pk.View1D[pk.double], out: pk.View1D[pk.double]):
+    out[tid] = atanh(view[tid])
+
+
+@pk.workunit
+def atanh_impl_2d_double(tid: int, view: pk.View2D[pk.double], out: pk.View2D[pk.double]):
+    for i in range(view.extent(1)):
+        out[tid][i] = atanh(view[tid][i])
+
+
+@pk.workunit
+def atanh_impl_1d_float(tid: int, view: pk.View1D[pk.float], out: pk.View1D[pk.float]):
+    out[tid] = atanh(view[tid])
+
+
+@pk.workunit
+def atanh_impl_2d_float(tid: int, view: pk.View2D[pk.float], out: pk.View2D[pk.float]):
+    for i in range(view.extent(1)):
+        out[tid][i] = atanh(view[tid][i])
+
+
 @pk.workunit
 def floor_impl_1d_double(tid: int, view: pk.View1D[pk.double], out: pk.View1D[pk.double]):
     out[tid] = floor(view[tid])

pykokkos/lib/ufuncs.py

@@ -2679,3 +2679,37 @@ def floor(view):
                              out=out,
                              view=view)
     return out
+
+
+def atanh(view):
+    """
+    Calculates an approximation to the inverse hyperbolic tangent for each element ``x_i`` of the input view.
+
+    Parameters
+    ----------
+    view : pykokkos view
+            Input view whose elements each represent the area of a hyperbolic sector. Should have a floating-point data type.
+
+    Returns
+    -------
+    y : pykokkos view
+        A view containing the inverse hyperbolic tangent of each element in the input view. The returned view
+        must have a floating-point data type determined by type promotion rules.
+    """
+    dtype = view.dtype
+    ndims = len(view.shape)
+    if ndims > 2:
+        raise NotImplementedError("atanh() ufunc only supports up to 2D views")
+    out = pk.View([*view.shape], dtype=dtype)
+    if view.shape == ():
+        tid = 1
+    else:
+        tid = view.shape[0]
+    _ufunc_kernel_dispatcher(tid=tid,
+                             dtype=dtype,
+                             ndims=ndims,
+                             op="atanh",
+                             sub_dispatcher=pk.parallel_for,
+                             out=out,
+                             view=view)
+    return out