add trace (#651)

* add trace

Author: Rubtsowa

dpnp/backend/include/dpnp_iface.hpp

@@ -705,6 +705,18 @@ INP_DLLEXPORT void dpnp_std_c(
 template <typename _DataType, typename _IndecesType>
 INP_DLLEXPORT void dpnp_take_c(void* array, void* indices, void* result, size_t size);
 
+/**
+ * @ingroup BACKEND_API
+ * @brief math library implementation of trace function
+ *
+ * @param [in]  array      Input array with data.
+ * @param [out] result     Output array.
+ * @param [in]  shape      Shape of input array.
+ * @param [in]  ndim       Number of elements in array.shape.
+ */
+template <typename _DataType, typename _ResultType>
+INP_DLLEXPORT void dpnp_trace_c(const void* array, void* result, const size_t* shape, const size_t ndim);
+
 /**
  * @ingroup BACKEND_API
  * @brief math library implementation of take function

dpnp/backend/include/dpnp_iface_fptr.hpp

@@ -199,6 +199,7 @@ enum class DPNPFuncName : size_t
     DPNP_FN_TAN,                      /**< Used in numpy.tan() implementation  */
     DPNP_FN_TANH,                     /**< Used in numpy.tanh() implementation  */
     DPNP_FN_TRANSPOSE,                /**< Used in numpy.transpose() implementation  */
+    DPNP_FN_TRACE,                    /**< Used in numpy.trace() implementation  */
     DPNP_FN_TRAPZ,                    /**< Used in numpy.trapz() implementation  */
     DPNP_FN_TRI,                      /**< Used in numpy.tri() implementation  */
     DPNP_FN_TRIL,                     /**< Used in numpy.tril() implementation  */

dpnp/backend/kernels/dpnp_krnl_arraycreation.cpp

@@ -211,6 +211,63 @@ void dpnp_vander_c(const void* array1_in, void* result1, const size_t size_in, c
     }
 }
 
+template <typename _DataType, typename _ResultType>
+class dpnp_trace_c_kernel;
+
+template <typename _DataType, typename _ResultType>
+void dpnp_trace_c(const void* array1_in, void* result1, const size_t* shape, const size_t ndim)
+{
+    cl::sycl::event event;
+
+    if ((array1_in == nullptr) || (result1 == nullptr))
+    {
+        return;
+    }
+
+    const _DataType* array_in = reinterpret_cast<const _DataType*>(array1_in);
+    _ResultType* result = reinterpret_cast<_ResultType*>(result1);
+
+    if (shape == nullptr)
+    {
+        return;
+    }
+
+    if (ndim == 0)
+    {
+        return;
+    }
+
+    size_t size = 1;
+    for (size_t i = 0; i < ndim - 1; ++i)
+    {
+        size *= shape[i];
+    }
+
+    if (size == 0)
+    {
+        return;
+    }
+
+    cl::sycl::range<1> gws(size);
+    auto kernel_parallel_for_func = [=](cl::sycl::id<1> global_id) {
+        size_t i = global_id[0];
+        _DataType elem = 0;
+        for (size_t j = 0; j < shape[ndim - 1]; ++j)
+        {
+            elem += array_in[i * shape[ndim - 1] + j];
+        }
+        result[i] = elem;
+    };
+
+    auto kernel_func = [&](cl::sycl::handler& cgh) {
+        cgh.parallel_for<class dpnp_trace_c_kernel<_DataType, _ResultType>>(gws, kernel_parallel_for_func);
+    };
+
+    event = DPNP_QUEUE.submit(kernel_func);
+
+    event.wait();
+}
+
 template <typename _DataType>
 class dpnp_tri_c_kernel;
 
@@ -539,6 +596,23 @@ void func_map_init_arraycreation(func_map_t& fmap)
     fmap[DPNPFuncName::DPNP_FN_VANDER][eft_C128][eft_C128] = {
         eft_C128, (void*)dpnp_vander_c<std::complex<double>, std::complex<double>>};
 
+    fmap[DPNPFuncName::DPNP_FN_TRACE][eft_INT][eft_INT] = {eft_INT, (void*)dpnp_trace_c<int, int>};
+    fmap[DPNPFuncName::DPNP_FN_TRACE][eft_LNG][eft_INT] = {eft_INT, (void*)dpnp_trace_c<long, int>};
+    fmap[DPNPFuncName::DPNP_FN_TRACE][eft_FLT][eft_INT] = {eft_INT, (void*)dpnp_trace_c<float, int>};
+    fmap[DPNPFuncName::DPNP_FN_TRACE][eft_DBL][eft_INT] = {eft_INT, (void*)dpnp_trace_c<double, int>};
+    fmap[DPNPFuncName::DPNP_FN_TRACE][eft_INT][eft_LNG] = {eft_LNG, (void*)dpnp_trace_c<int, long>};
+    fmap[DPNPFuncName::DPNP_FN_TRACE][eft_LNG][eft_LNG] = {eft_LNG, (void*)dpnp_trace_c<long, long>};
+    fmap[DPNPFuncName::DPNP_FN_TRACE][eft_FLT][eft_LNG] = {eft_LNG, (void*)dpnp_trace_c<float, long>};
+    fmap[DPNPFuncName::DPNP_FN_TRACE][eft_DBL][eft_LNG] = {eft_LNG, (void*)dpnp_trace_c<double, long>};
+    fmap[DPNPFuncName::DPNP_FN_TRACE][eft_INT][eft_FLT] = {eft_FLT, (void*)dpnp_trace_c<int, float>};
+    fmap[DPNPFuncName::DPNP_FN_TRACE][eft_LNG][eft_FLT] = {eft_FLT, (void*)dpnp_trace_c<long, float>};
+    fmap[DPNPFuncName::DPNP_FN_TRACE][eft_FLT][eft_FLT] = {eft_FLT, (void*)dpnp_trace_c<float, float>};
+    fmap[DPNPFuncName::DPNP_FN_TRACE][eft_DBL][eft_FLT] = {eft_FLT, (void*)dpnp_trace_c<double, float>};
+    fmap[DPNPFuncName::DPNP_FN_TRACE][eft_INT][eft_DBL] = {eft_DBL, (void*)dpnp_trace_c<int, double>};
+    fmap[DPNPFuncName::DPNP_FN_TRACE][eft_LNG][eft_DBL] = {eft_DBL, (void*)dpnp_trace_c<long, double>};
+    fmap[DPNPFuncName::DPNP_FN_TRACE][eft_FLT][eft_DBL] = {eft_DBL, (void*)dpnp_trace_c<float, double>};
+    fmap[DPNPFuncName::DPNP_FN_TRACE][eft_DBL][eft_DBL] = {eft_DBL, (void*)dpnp_trace_c<double, double>};
+
     fmap[DPNPFuncName::DPNP_FN_TRI][eft_INT][eft_INT] = {eft_INT, (void*)dpnp_tri_c<int>};
     fmap[DPNPFuncName::DPNP_FN_TRI][eft_LNG][eft_LNG] = {eft_LNG, (void*)dpnp_tri_c<long>};
     fmap[DPNPFuncName::DPNP_FN_TRI][eft_FLT][eft_FLT] = {eft_FLT, (void*)dpnp_tri_c<float>};

dpnp/dpnp_algo/dpnp_algo.pxd

@@ -171,6 +171,7 @@ cdef extern from "dpnp_iface_fptr.hpp" namespace "DPNPFuncName":  # need this na
         DPNP_FN_TAKE
         DPNP_FN_TAN
         DPNP_FN_TANH
+        DPNP_FN_TRACE
         DPNP_FN_TRANSPOSE
         DPNP_FN_TRAPZ
         DPNP_FN_TRI

dpnp/dpnp_algo/dpnp_algo_arraycreation.pyx

@@ -52,6 +52,7 @@ __all__ += [
     "dpnp_meshgrid",
     "dpnp_ones",
     "dpnp_ones_like",
+    "dpnp_trace",
     "dpnp_tri",
     "dpnp_tril",
     "dpnp_triu",
@@ -64,6 +65,7 @@ __all__ += [
 ctypedef void(*custom_1in_1out_func_ptr_t)(void *, void * , const int , size_t * , size_t * , const size_t, const size_t)
 ctypedef void(*ftpr_custom_vander_1in_1out_t)(void *, void *, size_t, size_t, int)
 ctypedef void(*custom_indexing_1out_func_ptr_t)(void * , const size_t , const size_t , const int)
+ctypedef void(*fptr_dpnp_trace_t)(const void *, void * , const size_t * , const size_t)
 
 
 cpdef dparray dpnp_copy(dparray x1, order, subok):
@@ -270,6 +272,29 @@ cpdef dparray dpnp_ones_like(result_shape, result_dtype):
     return call_fptr_1out(DPNP_FN_ONES_LIKE, result_shape, result_dtype)
 
 
+cpdef dparray dpnp_trace(arr, offset=0, axis1=0, axis2=1, dtype=None, out=None):
+    if dtype is None:
+        dtype_ = arr.dtype
+    else:
+        dtype_ = dtype
+
+    cdef dparray diagonal_arr = dpnp.diagonal(arr, offset, axis1, axis2)
+
+    cdef DPNPFuncType param1_type = dpnp_dtype_to_DPNPFuncType(arr.dtype)
+    cdef DPNPFuncType param2_type = dpnp_dtype_to_DPNPFuncType(dtype_)
+
+    cdef DPNPFuncData kernel_data = get_dpnp_function_ptr(DPNP_FN_TRACE, param1_type, param2_type)
+
+    result_type = dpnp_DPNPFuncType_to_dtype( < size_t > kernel_data.return_type)
+    cdef dparray result = dparray(diagonal_arr.shape[:-1], dtype=result_type)
+
+    cdef fptr_dpnp_trace_t func = <fptr_dpnp_trace_t > kernel_data.ptr
+
+    func(diagonal_arr.get_data(), result.get_data(), < size_t * > diagonal_arr._dparray_shape.data(), diagonal_arr.ndim)
+
+    return result
+
+
 cpdef dparray dpnp_tri(N, M=None, k=0, dtype=numpy.float):
     if M is None:
         M = N

dpnp/dpnp_iface_arraycreation.py

@@ -75,6 +75,7 @@
     "ogrid",
     "ones",
     "ones_like",
+    "trace",
     "tri",
     "tril",
     "triu",
@@ -1072,6 +1073,36 @@ def ones_like(x1, dtype=None, order='C', subok=False, shape=None):
     return numpy.ones_like(x1, dtype, order, subok, shape)
 
 
+def trace(arr, offset=0, axis1=0, axis2=1, dtype=None, out=None):
+    """
+       Return the sum along diagonals of the array.
+
+       For full documentation refer to :obj:`numpy.trace`.
+
+       Limitations
+       -----------
+       Input array is supported as :obj:`dpnp.ndarray`.
+       Parameters ``axis1``, ``axis2``, ``out`` and ``dtype`` are supported only with default values.
+       """
+    if not use_origin_backend():
+        if not isinstance(arr, dparray):
+            pass
+        elif arr.size == 0:
+            pass
+        elif arr.ndim < 2:
+            pass
+        elif axis1 != 0:
+            pass
+        elif axis2 != 1:
+            pass
+        elif out is not None and (not isinstance(out, dparray) or (isinstance(out, dparray) and out.shape != arr.shape)):
+            pass
+        else:
+            return dpnp_trace(arr, offset, axis1, axis2, dtype, out)
+
+    return call_origin(numpy.trace, arr, offset, axis1, axis2, dtype, out)
+
+
 def tri(N, M=None, k=0, dtype=numpy.float, **kwargs):
     """
     An array with ones at and below the given diagonal and zeros elsewhere.

tests/test_arraycreation.py

@@ -151,6 +151,43 @@ def test_loadtxt(type):
         numpy.testing.assert_array_equal(dpnp_res, np_res)
 
 
+@pytest.mark.parametrize("dtype",
+                         [numpy.float64, numpy.float32, numpy.int64, numpy.int32],
+                         ids=['float64', 'float32', 'int64', 'int32'])
+@pytest.mark.parametrize("type",
+                         [numpy.float64, numpy.float32, numpy.int64, numpy.int32],
+                         ids=['float64', 'float32', 'int64', 'int32'])
+@pytest.mark.parametrize("offset",
+                         [0, 1],
+                         ids=['0', '1'])
+@pytest.mark.parametrize("array",
+                         [[[0, 0], [0, 0]],
+                          [[1, 2], [1, 2]],
+                          [[1, 2], [3, 4]],
+                          [[0, 1, 2], [3, 4, 5], [6, 7, 8]],
+                          [[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]],
+                          [[[1, 2], [3, 4]], [[1, 2], [2, 1]], [[1, 3], [3, 1]]],
+                          [[[[1, 2], [3, 4]], [[1, 2], [2, 1]]], [[[1, 3], [3, 1]], [[0, 1], [1, 3]]]],
+                          [[[[1, 2, 3], [3, 4, 5]], [[1, 2, 3], [2, 1, 0]]], [
+                              [[1, 3, 5], [3, 1, 0]], [[0, 1, 2], [1, 3, 4]]]],
+                          [[[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]], [[[13, 14, 15], [16, 17, 18]], [[19, 20, 21], [22, 23, 24]]]]],
+                         ids=['[[0, 0], [0, 0]]',
+                              '[[1, 2], [1, 2]]',
+                              '[[1, 2], [3, 4]]',
+                              '[[0, 1, 2], [3, 4, 5], [6, 7, 8]]',
+                              '[[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]',
+                              '[[[1, 2], [3, 4]], [[1, 2], [2, 1]], [[1, 3], [3, 1]]]',
+                              '[[[[1, 2], [3, 4]], [[1, 2], [2, 1]]], [[[1, 3], [3, 1]], [[0, 1], [1, 3]]]]',
+                              '[[[[1, 2, 3], [3, 4, 5]], [[1, 2, 3], [2, 1, 0]]], [[[1, 3, 5], [3, 1, 0]], [[0, 1, 2], [1, 3, 4]]]]',
+                              '[[[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]], [[[13, 14, 15], [16, 17, 18]], [[19, 20, 21], [22, 23, 24]]]]'])
+def test_trace(array, offset, type, dtype):
+    a = numpy.array(array, type)
+    ia = dpnp.array(array, type)
+    expected = numpy.trace(a, offset=offset, dtype=dtype)
+    result = dpnp.trace(ia, offset=offset, dtype=dtype)
+    numpy.testing.assert_array_equal(expected, result)
+
+
 @pytest.mark.parametrize("N",
                          [0, 1, 2, 3, 4],
                          ids=['0', '1', '2', '3', '4'])

tests_external/skipped_tests_numpy.tbl

@@ -1145,7 +1145,6 @@ tests/test_numeric.py::TestNonarrayArgs::test_reshape
 tests/test_numeric.py::TestNonarrayArgs::test_searchsorted
 tests/test_numeric.py::TestNonarrayArgs::test_size
 tests/test_numeric.py::TestNonarrayArgs::test_std
-tests/test_numeric.py::TestNonarrayArgs::test_trace
 tests/test_numeric.py::TestNonarrayArgs::test_var
 tests/test_numeric.py::TestNonzero::test_array_method
 tests/test_numeric.py::TestNonzero::test_count_nonzero_axis