diff --git a/.github/workflows/array_api.yml b/.github/workflows/array_api.yml index 2d656c21..8e762d03 100644 --- a/.github/workflows/array_api.yml +++ b/.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 diff --git a/pykokkos/__init__.py b/pykokkos/__init__.py index 3b83a079..102276d6 100644 --- a/pykokkos/__init__.py +++ b/pykokkos/__init__.py @@ -38,6 +38,7 @@ sin, cos, tan, + atanh, logical_and, logical_or, logical_xor, diff --git a/pykokkos/lib/ufunc_workunits.py b/pykokkos/lib/ufunc_workunits.py index 06f12530..81b5a28b 100644 --- a/pykokkos/lib/ufunc_workunits.py +++ b/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]) diff --git a/pykokkos/lib/ufuncs.py b/pykokkos/lib/ufuncs.py index 13be86e2..95bbc86d 100644 --- a/pykokkos/lib/ufuncs.py +++ b/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