add identity function (#652)

KsanaKozlova · web-flow · commit 1b4ab46c23ae · 2021-03-26T14:03:20.000-05:00
* add identity function
diff --git a/dpnp/backend/include/dpnp_iface.hpp b/dpnp/backend/include/dpnp_iface.hpp
@@ -485,6 +485,16 @@ template <typename _DataType>
 INP_DLLEXPORT void dpnp_diagonal_c(
     void* array1_in, void* result1, const size_t offset, size_t* shape, size_t* res_shape, const size_t res_ndim);
 
+/**
+ * @ingroup BACKEND_API
+ * @brief Implementation of identity function
+ *
+ * @param [out] result1   Output array.
+ * @param [in]  n         Number of rows (and columns) in n x n output.
+ */
+template <typename _DataType>
+INP_DLLEXPORT void dpnp_identity_c(void* result1, const size_t n);
+
 /**
  * @ingroup BACKEND_API
  * @brief implementation of creating filled with value array function
diff --git a/dpnp/backend/include/dpnp_iface_fptr.hpp b/dpnp/backend/include/dpnp_iface_fptr.hpp
@@ -114,6 +114,7 @@ enum class DPNPFuncName : size_t
     DPNP_FN_FULL,                     /**< Used in numpy.full() implementation  */
     DPNP_FN_FULL_LIKE,                /**< Used in numpy.full_like() implementation  */
     DPNP_FN_HYPOT,                    /**< Used in numpy.hypot() implementation  */
+    DPNP_FN_IDENTITY,                 /**< Used in numpy.identity() implementation  */
     DPNP_FN_INITVAL,                  /**< Used in numpy ones, ones_like, zeros, zeros_like, full, full_like impl  */
     DPNP_FN_INV,                      /**< Used in numpy.linalg.inv() implementation  */
     DPNP_FN_INVERT,                   /**< Used in numpy.invert() implementation  */
diff --git a/dpnp/backend/kernels/dpnp_krnl_arraycreation.cpp b/dpnp/backend/kernels/dpnp_krnl_arraycreation.cpp
@@ -110,6 +110,37 @@ void dpnp_full_like_c(void* array_in, void* result, const size_t size)
     dpnp_full_c<_DataType>(array_in, result, size);
 }
 
+template <typename _KernelNameSpecialization>
+class dpnp_identity_c_kernel;
+
+template <typename _DataType>
+void dpnp_identity_c(void* result1, const size_t n)
+{
+    if (n == 0)
+    {
+        return;
+    }
+
+    cl::sycl::event event;
+
+    _DataType* result = reinterpret_cast<_DataType*>(result1);
+
+    cl::sycl::range<2> gws(n, n);
+    auto kernel_parallel_for_func = [=](cl::sycl::id<2> global_id) {
+        size_t i = global_id[0];
+        size_t j = global_id[1];
+        result[i * n + j] = i == j;
+    };
+
+    auto kernel_func = [&](cl::sycl::handler& cgh) {
+        cgh.parallel_for<class dpnp_identity_c_kernel<_DataType>>(gws, kernel_parallel_for_func);
+    };
+
+    event = DPNP_QUEUE.submit(kernel_func);
+
+    event.wait();  
+}
+
 template <typename _DataType>
 void dpnp_ones_c(void* result, size_t size)
 {
@@ -478,6 +509,13 @@ void func_map_init_arraycreation(func_map_t& fmap)
     fmap[DPNPFuncName::DPNP_FN_FULL_LIKE][eft_C128][eft_C128] = {eft_C128,
                                                                  (void*)dpnp_full_like_c<std::complex<double>>};
 
+    fmap[DPNPFuncName::DPNP_FN_IDENTITY][eft_INT][eft_INT] = {eft_INT, (void*)dpnp_identity_c<int>};
+    fmap[DPNPFuncName::DPNP_FN_IDENTITY][eft_LNG][eft_LNG] = {eft_LNG, (void*)dpnp_identity_c<long>};
+    fmap[DPNPFuncName::DPNP_FN_IDENTITY][eft_FLT][eft_FLT] = {eft_FLT, (void*)dpnp_identity_c<float>};
+    fmap[DPNPFuncName::DPNP_FN_IDENTITY][eft_DBL][eft_DBL] = {eft_DBL, (void*)dpnp_identity_c<double>};
+    fmap[DPNPFuncName::DPNP_FN_IDENTITY][eft_BLN][eft_BLN] = {eft_BLN, (void*)dpnp_identity_c<bool>};
+    fmap[DPNPFuncName::DPNP_FN_IDENTITY][eft_C128][eft_C128] = {eft_C128, (void*)dpnp_identity_c<std::complex<double>>};
+
     fmap[DPNPFuncName::DPNP_FN_ONES][eft_INT][eft_INT] = {eft_INT, (void*)dpnp_ones_c<int>};
     fmap[DPNPFuncName::DPNP_FN_ONES][eft_LNG][eft_LNG] = {eft_LNG, (void*)dpnp_ones_c<long>};
     fmap[DPNPFuncName::DPNP_FN_ONES][eft_FLT][eft_FLT] = {eft_FLT, (void*)dpnp_ones_c<float>};
diff --git a/dpnp/dpnp_algo/dpnp_algo.pxd b/dpnp/dpnp_algo/dpnp_algo.pxd
@@ -87,6 +87,7 @@ cdef extern from "dpnp_iface_fptr.hpp" namespace "DPNPFuncName":  # need this na
         DPNP_FN_FULL
         DPNP_FN_FULL_LIKE
         DPNP_FN_HYPOT
+        DPNP_FN_IDENTITY
         DPNP_FN_INITVAL
         DPNP_FN_INV
         DPNP_FN_INVERT
diff --git a/dpnp/dpnp_algo/dpnp_algo_arraycreation.pyx b/dpnp/dpnp_algo/dpnp_algo_arraycreation.pyx
@@ -46,6 +46,7 @@ __all__ += [
     "dpnp_full",
     "dpnp_full_like",
     "dpnp_geomspace",
+    "dpnp_identity",
     "dpnp_linspace",
     "dpnp_logspace",
     "dpnp_meshgrid",
@@ -167,6 +168,21 @@ cpdef dparray dpnp_geomspace(start, stop, num, endpoint, dtype, axis):
     return result
 
 
+cpdef dparray dpnp_identity(n, result_dtype):
+    cdef DPNPFuncType dtype_in = dpnp_dtype_to_DPNPFuncType(result_dtype)
+
+    cdef DPNPFuncData kernel_data = get_dpnp_function_ptr(DPNP_FN_IDENTITY, dtype_in, DPNP_FT_NONE)
+
+    result_type = dpnp_DPNPFuncType_to_dtype( < size_t > kernel_data.return_type)
+
+    cdef dparray result = dparray((n, n), dtype=result_type)
+
+    cdef fptr_1out_t func = <fptr_1out_t > kernel_data.ptr
+    func(result.get_data(), n)
+
+    return result
+
+
 # TODO this function should work through dpnp_arange_c
 cpdef tuple dpnp_linspace(start, stop, num, endpoint, retstep, dtype, axis):
     cdef dparray result = dparray(num, dtype=dtype)
diff --git a/dpnp/dpnp_iface_arraycreation.py b/dpnp/dpnp_iface_arraycreation.py
@@ -66,6 +66,7 @@
     "full",
     "full_like",
     "geomspace",
+    "identity",
     "linspace",
     "loadtxt",
     "logspace",
@@ -711,6 +712,40 @@ def geomspace(start, stop, num=50, endpoint=True, dtype=None, axis=0):
     return call_origin(numpy.geomspace, start, stop, num, endpoint, dtype, axis)
 
 
+def identity(n, dtype=None, *, like=None):
+    """
+    Return the identity array.
+
+    The identity array is a square array with ones on the main diagonal.
+
+    For full documentation refer to :obj:`numpy.identity`.
+
+    Limitations
+    -----------
+    Parameter ``like`` is currently not supported .
+
+    Examples
+    --------
+    >>> import dpnp as np
+    >>> np.identity(3)
+    array([[1.,  0.,  0.],
+           [0.,  1.,  0.],
+           [0.,  0.,  1.]])
+
+    """
+    if not use_origin_backend():
+        if like is not None:
+            pass
+        elif n < 0:
+            pass
+        else:
+            if dtype is None:
+                dtype = dpnp.float64
+            return dpnp_identity(n, dtype)
+
+    return call_origin(numpy.identity, n, dtype=dtype, like=like)
+
+
 def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None, axis=0):
     """
     Return evenly spaced numbers over a specified interval.
diff --git a/tests/skipped_tests.tbl b/tests/skipped_tests.tbl
@@ -325,7 +325,6 @@ tests/third_party/cupy/creation_tests/test_basic.py::TestBasic::test_eye
 tests/third_party/cupy/creation_tests/test_basic.py::TestBasic::test_full_default_dtype
 tests/third_party/cupy/creation_tests/test_basic.py::TestBasic::test_full_default_dtype_cpu_input
 tests/third_party/cupy/creation_tests/test_basic.py::TestBasic::test_full_like_subok
-tests/third_party/cupy/creation_tests/test_basic.py::TestBasic::test_identity
 tests/third_party/cupy/creation_tests/test_basic.py::TestBasic::test_ones_like_subok
 tests/third_party/cupy/creation_tests/test_basic.py::TestBasic::test_zeros_like_subok
 tests/third_party/cupy/creation_tests/test_basic.py::TestBasic::test_zeros_strides
diff --git a/tests/skipped_tests_gpu.tbl b/tests/skipped_tests_gpu.tbl
@@ -389,7 +389,6 @@ tests/third_party/cupy/creation_tests/test_basic.py::TestBasic::test_eye
 tests/third_party/cupy/creation_tests/test_basic.py::TestBasic::test_full_default_dtype
 tests/third_party/cupy/creation_tests/test_basic.py::TestBasic::test_full_default_dtype_cpu_input
 tests/third_party/cupy/creation_tests/test_basic.py::TestBasic::test_full_like_subok
-tests/third_party/cupy/creation_tests/test_basic.py::TestBasic::test_identity
 tests/third_party/cupy/creation_tests/test_basic.py::TestBasic::test_ones_like_subok
 tests/third_party/cupy/creation_tests/test_basic.py::TestBasic::test_zeros_like_subok
 tests/third_party/cupy/creation_tests/test_basic.py::TestBasic::test_zeros_strides
diff --git a/tests/test_arraycreation.py b/tests/test_arraycreation.py
@@ -122,6 +122,19 @@ def test_geomspace(type, num, endpoint):
         numpy.testing.assert_allclose(dpnp_res, np_res)
 
 
+@pytest.mark.parametrize("n",
+                        [0, 1, 4],
+                        ids=['0', '1', '4'])
+@pytest.mark.parametrize("type",
+                        [numpy.float64, numpy.float32, numpy.int64,
+                        numpy.int32, numpy.bool, numpy.complex128, None],
+                        ids=['float64', 'float32', 'int64', 'int32', 'bool', 'complex128', 'None'])
+def test_identity(n, type):
+    expected = numpy.identity(n, dtype=type)
+    result = dpnp.identity(n, dtype=type)
+    numpy.testing.assert_array_equal(expected, result)
+
+
 @pytest.mark.parametrize("type",
                          [numpy.float64, numpy.float32, numpy.int64, numpy.int32],
                          ids=['float64', 'float32', 'int64', 'int32'])
