Move ti.imag()/isfinite()/isfinite() and reuse them

Author: vlad-perevezentsev

dpctl_ext/tensor/CMakeLists.txt

@@ -103,9 +103,9 @@ set(_elementwise_sources
     #${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/elementwise_functions/greater_equal.cpp
     #${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/elementwise_functions/greater.cpp
     #${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/elementwise_functions/hypot.cpp
-    #${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/elementwise_functions/imag.cpp
-    #${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/elementwise_functions/isfinite.cpp
-    #${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/elementwise_functions/isinf.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/elementwise_functions/imag.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/elementwise_functions/isfinite.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/elementwise_functions/isinf.cpp
     #${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/elementwise_functions/isnan.cpp
     #${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/elementwise_functions/less_equal.cpp
     #${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/elementwise_functions/less.cpp

dpctl_ext/tensor/__init__.py

@@ -100,6 +100,9 @@
     exp,
     expm1,
     floor,
+    imag,
+    isfinite,
+    isinf,
 )
 from ._reduction import (
     argmax,
@@ -171,6 +174,9 @@
     "full",
     "full_like",
     "iinfo",
+    "imag",
+    "isfinite",
+    "isinf",
     "isdtype",
     "isin",
     "linspace",

dpctl_ext/tensor/_elementwise_funcs.py

@@ -464,6 +464,94 @@
 )
 del _floor_docstring
 
+# U16: ==== IMAG        (x)
+_imag_docstring = r"""
+imag(x, /, \*, out=None, order='K')
+
+Computes imaginary part of each element `x_i` for input array `x`.
+
+Args:
+    x (usm_ndarray):
+        Input array. May have any data type.
+    out (Union[usm_ndarray, None], optional):
+        Output array to populate.
+        Array must have the correct shape and the expected data type.
+    order ("C","F","A","K", optional):
+        Memory layout of the new output array, if parameter
+        `out` is ``None``.
+        Default: "K".
+
+Returns:
+    usm_ndarray:
+        An array containing the element-wise imaginary component of input.
+        If the input is a real-valued data type, the returned array has
+        the same data type. If the input is a complex floating-point
+        data type, the returned array has a floating-point data type
+        with the same floating-point precision as complex input.
+"""
+
+imag = UnaryElementwiseFunc(
+    "imag", ti._imag_result_type, ti._imag, _imag_docstring
+)
+del _imag_docstring
+
+# U17: ==== ISFINITE    (x)
+_isfinite_docstring_ = r"""
+isfinite(x, /, \*, out=None, order='K')
+
+Test if each element of input array is a finite number.
+
+Args:
+    x (usm_ndarray):
+        Input array. May have any data type.
+    out (Union[usm_ndarray, None], optional):
+        Output array to populate.
+        Array must have the correct shape and the expected data type.
+    order ("C","F","A","K", optional):
+        Memory layout of the new output array, if parameter
+        `out` is ``None``.
+        Default: "K".
+
+Returns:
+    usm_ndarray:
+        An array which is True where `x` is not positive infinity,
+        negative infinity, or NaN, False otherwise.
+        The data type of the returned array is `bool`.
+"""
+
+isfinite = UnaryElementwiseFunc(
+    "isfinite", ti._isfinite_result_type, ti._isfinite, _isfinite_docstring_
+)
+del _isfinite_docstring_
+
+# U18: ==== ISINF       (x)
+_isinf_docstring_ = r"""
+isinf(x, /, \*, out=None, order='K')
+
+Test if each element of input array is an infinity.
+
+Args:
+    x (usm_ndarray):
+        Input array. May have any data type.
+    out (Union[usm_ndarray, None], optional):
+        Output array to populate.
+        Array must have the correct shape and the expected data type.
+    order ("C","F","A","K", optional):
+        Memory layout of the new output array, if parameter
+        `out` is ``None``.
+        Default: "K".
+
+Returns:
+    usm_ndarray:
+        An array which is True where `x` is positive or negative infinity,
+        False otherwise. The data type of the returned array is `bool`.
+"""
+
+isinf = UnaryElementwiseFunc(
+    "isinf", ti._isinf_result_type, ti._isinf, _isinf_docstring_
+)
+del _isinf_docstring_
+
 # U43: ==== ANGLE        (x)
 _angle_docstring = r"""
 angle(x, /, \*, out=None, order='K')

dpctl_ext/tensor/libtensor/include/kernels/elementwise_functions/exp.hpp

@@ -26,16 +26,20 @@
 // THE POSSIBILITY OF SUCH DAMAGE.
 //*****************************************************************************
 //
+//===---------------------------------------------------------------------===//
+///
 /// \file
 /// This file defines kernels for elementwise evaluation of EXP(x) function.
 //===---------------------------------------------------------------------===//
 
 #pragma once
 #include <cmath>
+#include <complex>
 #include <cstddef>
 #include <cstdint>
 #include <limits>
 #include <type_traits>
+#include <vector>
 
 #include <sycl/sycl.hpp>
 
@@ -45,7 +49,6 @@
 #include "kernels/dpctl_tensor_types.hpp"
 #include "kernels/elementwise_functions/common.hpp"
 
-#include "utils/offset_utils.hpp"
 #include "utils/type_dispatch_building.hpp"
 #include "utils/type_utils.hpp"
 

dpctl_ext/tensor/libtensor/include/kernels/elementwise_functions/expm1.hpp

@@ -26,16 +26,20 @@
 // THE POSSIBILITY OF SUCH DAMAGE.
 //*****************************************************************************
 //
+//===---------------------------------------------------------------------===//
+///
 /// \file
 /// This file defines kernels for elementwise evaluation of EXPM1(x) function.
 //===---------------------------------------------------------------------===//
 
 #pragma once
 #include <cmath>
+#include <complex>
 #include <cstddef>
 #include <cstdint>
 #include <limits>
 #include <type_traits>
+#include <vector>
 
 #include <sycl/sycl.hpp>
 
@@ -44,7 +48,6 @@
 #include "kernels/dpctl_tensor_types.hpp"
 #include "kernels/elementwise_functions/common.hpp"
 
-#include "utils/offset_utils.hpp"
 #include "utils/type_dispatch_building.hpp"
 #include "utils/type_utils.hpp"
 

dpctl_ext/tensor/libtensor/include/kernels/elementwise_functions/floor.hpp

@@ -26,6 +26,8 @@
 // THE POSSIBILITY OF SUCH DAMAGE.
 //*****************************************************************************
 //
+//===---------------------------------------------------------------------===//
+///
 /// \file
 /// This file defines kernels for elementwise evaluation of FLOOR(x) function.
 //===---------------------------------------------------------------------===//
@@ -35,6 +37,7 @@
 #include <cstddef>
 #include <cstdint>
 #include <type_traits>
+#include <vector>
 
 #include <sycl/sycl.hpp>
 
@@ -43,7 +46,6 @@
 #include "kernels/dpctl_tensor_types.hpp"
 #include "kernels/elementwise_functions/common.hpp"
 
-#include "utils/offset_utils.hpp"
 #include "utils/type_dispatch_building.hpp"
 #include "utils/type_utils.hpp"