Intexing mod to desc part1 (#778)

shssf · web-flow · commit ff83dbae60aa · 2021-07-11T12:42:23.000-05:00
diff --git a/dpnp/dpnp_algo/dpnp_algo_indexing.pyx b/dpnp/dpnp_algo/dpnp_algo_indexing.pyx
@@ -202,7 +202,7 @@ cpdef dpnp_place(dparray arr, dparray mask, dparray vals):
     func(arr.get_data(), mask_.get_data(), vals.get_data(), arr.size, vals.size)
 
 
-cpdef dpnp_put(dparray input, ind, v):
+cpdef dpnp_put(dpnp_descriptor input, object ind, v):
     ind_is_list = isinstance(ind, list)
 
     if dpnp.isscalar(ind):
@@ -236,17 +236,18 @@ cpdef dpnp_put(dparray input, ind, v):
     func(input.get_data(), ind_array.get_data(), v_array.get_data(), input.size, ind_array.size, v_array.size)
 
 
-cpdef dpnp_put_along_axis(dparray arr, dparray indices, dparray values, int axis):
+cpdef dpnp_put_along_axis(dpnp_descriptor arr, dpnp_descriptor indices, dpnp_descriptor values, int axis):
+    cdef dparray_shape_type arr_shape = arr.shape
     cdef DPNPFuncType param1_type = dpnp_dtype_to_DPNPFuncType(arr.dtype)
 
     cdef DPNPFuncData kernel_data = get_dpnp_function_ptr(DPNP_FN_PUT_ALONG_AXIS, param1_type, param1_type)
 
     cdef custom_indexing_3in_with_axis_func_ptr_t func = <custom_indexing_3in_with_axis_func_ptr_t > kernel_data.ptr
 
-    func(arr.get_data(), indices.get_data(), values.get_data(), axis, < size_t * > arr._dparray_shape.data(), arr.ndim, indices.size, values.size)
+    func(arr.get_data(), indices.get_data(), values.get_data(), axis, < size_t * > arr_shape.data(), arr.ndim, indices.size, values.size)
 
 
-cpdef dpnp_putmask(dparray arr, dparray mask, dparray values):
+cpdef dpnp_putmask(object arr, object mask, object values):
     cdef int values_size = values.size
     for i in range(arr.size):
         if mask[i]:
@@ -269,7 +270,7 @@ cpdef dparray dpnp_select(condlist, choicelist, default):
     return res_array.reshape(condlist[0].shape)
 
 
-cpdef dparray dpnp_take(dparray input, dparray indices):
+cpdef dparray dpnp_take(dpnp_descriptor input, dpnp_descriptor indices):
     cdef DPNPFuncType param1_type = dpnp_dtype_to_DPNPFuncType(input.dtype)
 
     cdef DPNPFuncData kernel_data = get_dpnp_function_ptr(DPNP_FN_TAKE, param1_type, param1_type)
@@ -284,7 +285,7 @@ cpdef dparray dpnp_take(dparray input, dparray indices):
     return result
 
 
-cpdef dparray dpnp_take_along_axis(dparray arr, dparray indices, int axis):
+cpdef dparray dpnp_take_along_axis(object arr, object indices, int axis):
     cdef long size_arr = arr.size
     cdef dparray_shape_type shape_arr = arr.shape
     cdef long size_indices = indices.size
@@ -391,7 +392,7 @@ cpdef tuple dpnp_tril_indices(n, k=0, m=None):
     return (dparray1, dparray2)
 
 
-cpdef tuple dpnp_tril_indices_from(arr, k=0):
+cpdef tuple dpnp_tril_indices_from(dpnp_descriptor arr, k=0):
     m = arr.shape[0]
     n = arr.shape[1]
     array1 = []
@@ -435,7 +436,7 @@ cpdef tuple dpnp_triu_indices(n, k=0, m=None):
     return (dparray1, dparray2)
 
 
-cpdef tuple dpnp_triu_indices_from(arr, k=0):
+cpdef tuple dpnp_triu_indices_from(dpnp_descriptor arr, k=0):
     m = arr.shape[0]
     n = arr.shape[1]
     array1 = []
diff --git a/dpnp/dpnp_iface_indexing.py b/dpnp/dpnp_iface_indexing.py
@@ -386,7 +386,7 @@ def place(arr, mask, vals):
     return call_origin(numpy.place, arr, mask, vals)
 
 
-def put(input, ind, v, mode='raise'):
+def put(x1, ind, v, mode='raise'):
     """
     Replaces specified elements of an array with given values.
     For full documentation refer to :obj:`numpy.put`.
@@ -397,22 +397,21 @@ def put(input, ind, v, mode='raise'):
     Not supported parameter mode.
     """
 
-    if not use_origin_backend(input):
-        if not isinstance(input, dparray):
-            pass
-        elif mode != 'raise':
+    x1_desc = dpnp.get_dpnp_descriptor(x1)
+    if x1_desc:
+        if mode != 'raise':
             pass
         elif type(ind) != type(v):
             pass
-        elif numpy.max(ind) >= input.size or numpy.min(ind) + input.size < 0:
+        elif numpy.max(ind) >= x1_desc.size or numpy.min(ind) + x1_desc.size < 0:
             pass
         else:
-            return dpnp_put(input, ind, v)
+            return dpnp_put(x1_desc, ind, v)
 
-    return call_origin(numpy.put, input, ind, v, mode)
+    return call_origin(numpy.put, x1, ind, v, mode)
 
 
-def put_along_axis(arr, indices, values, axis):
+def put_along_axis(x1, indices, values, axis):
     """
     Put values into the destination array by matching 1d index and data slices.
     For full documentation refer to :obj:`numpy.put_along_axis`.
@@ -422,62 +421,25 @@ def put_along_axis(arr, indices, values, axis):
     :obj:`take_along_axis` : Take values from the input array by matching 1d index and data slices.
     """
 
-    if not use_origin_backend(arr):
-        if not isinstance(arr, dparray):
-            pass
-        elif not isinstance(indices, dparray):
-            pass
-        elif arr.ndim != indices.ndim:
+    x1_desc = dpnp.get_dpnp_descriptor(x1)
+    indices_desc = dpnp.get_dpnp_descriptor(indices)
+    values_desc = dpnp.get_dpnp_descriptor(values)
+    if x1_desc and indices_desc and values_desc:
+        if x1_desc.ndim != indices_desc.ndim:
             pass
         elif not isinstance(axis, int):
             pass
-        elif axis >= arr.ndim:
+        elif axis >= x1_desc.ndim:
             pass
-        elif not isinstance(values, (dparray, tuple, list)) and not dpnp.isscalar(values):
+        elif indices_desc.size != values_desc.size:
             pass
-        elif not dpnp.isscalar(values) and ((isinstance(values, dparray) and indices.size != values.size) or
-                                            ((isinstance(values, (tuple, list)) and indices.size != len(values)))):
-            pass
-        elif arr.ndim == indices.ndim:
-            val_list = []
-            for i in list(indices.shape)[:-1]:
-                if i == 1:
-                    val_list.append(True)
-                else:
-                    val_list.append(False)
-            if not all(val_list):
-                pass
-            else:
-                if dpnp.isscalar(values):
-                    values_size = 1
-                    values_ = dparray(values_size, dtype=arr.dtype)
-                    values_[0] = values
-                elif isinstance(values, dparray):
-                    values_ = values
-                else:
-                    values_size = len(values)
-                    values_ = dparray(values_size, dtype=arr.dtype)
-                    for i in range(values_size):
-                        values_[i] = values[i]
-                return dpnp_put_along_axis(arr, indices, values_, axis)
         else:
-            if dpnp.isscalar(values):
-                values_size = 1
-                values_ = dparray(values_size, dtype=arr.dtype)
-                values_[0] = values
-            elif isinstance(values, dparray):
-                values_ = values
-            else:
-                values_size = len(values)
-                values_ = dparray(values_size, dtype=arr.dtype)
-                for i in range(values_size):
-                    values_[i] = values[i]
-            return dpnp_put_along_axis(arr, indices, values_, axis)
+            return dpnp_put_along_axis(x1_desc, indices_desc, values_desc, axis)
 
-    return call_origin(numpy.put_along_axis, arr, indices, values, axis)
+    return call_origin(numpy.put_along_axis, x1, indices, values, axis)
 
 
-def putmask(arr, mask, values):
+def putmask(x1, mask, values):
     """
     Changes elements of an array based on conditional and input values.
     For full documentation refer to :obj:`numpy.putmask`.
@@ -487,17 +449,13 @@ def putmask(arr, mask, values):
     Input arrays ``arr``, ``mask`` and ``values``  are supported as :obj:`dpnp.ndarray`.
     """
 
-    if not use_origin_backend(arr):
-        if not isinstance(arr, dparray):
-            pass
-        elif not isinstance(mask, dparray):
-            pass
-        elif not isinstance(values, dparray):
-            pass
-        else:
-            return dpnp_putmask(arr, mask, values)
+    x1_desc = dpnp.get_dpnp_descriptor(x1)
+    mask_desc = dpnp.get_dpnp_descriptor(mask)
+    values_desc = dpnp.get_dpnp_descriptor(values)
+    if x1_desc and mask_desc and values_desc:
+        return dpnp_putmask(x1, mask, values)
 
-    return call_origin(numpy.putmask, arr, mask, values)
+    return call_origin(numpy.putmask, x1, mask, values)
 
 
 def select(condlist, choicelist, default=0):
@@ -510,6 +468,7 @@ def select(condlist, choicelist, default=0):
     Arrays of input lists are supported as :obj:`dpnp.ndarray`.
     Parameter ``default`` are supported only with default values.
     """
+
     if not use_origin_backend():
         if not isinstance(condlist, list):
             pass
@@ -537,7 +496,7 @@ def select(condlist, choicelist, default=0):
     return call_origin(numpy.select, condlist, choicelist, default)
 
 
-def take(input, indices, axis=None, out=None, mode='raise'):
+def take(x1, indices, axis=None, out=None, mode='raise'):
     """
     Take elements from an array.
     For full documentation refer to :obj:`numpy.take`.
@@ -554,24 +513,22 @@ def take(input, indices, axis=None, out=None, mode='raise'):
     :obj:`take_along_axis` : Take elements by matching the array and the index arrays.
     """
 
-    if not use_origin_backend(input):
-        if not isinstance(input, dparray):
-            pass
-        elif not isinstance(indices, dparray):
-            pass
-        elif axis is not None:
+    x1_desc = dpnp.get_dpnp_descriptor(x1)
+    indices_desc = dpnp.get_dpnp_descriptor(indices)
+    if x1_desc and indices_desc:
+        if axis is not None:
             pass
         elif out is not None:
             pass
         elif mode != 'raise':
             pass
         else:
-            return dpnp_take(input, indices)
+            return dpnp_take(x1_desc, indices_desc)
 
-    return call_origin(numpy.take, input, indices, axis, out, mode)
+    return call_origin(numpy.take, x1, indices, axis, out, mode)
 
 
-def take_along_axis(arr, indices, axis):
+def take_along_axis(x1, indices, axis):
     """
     Take values from the input array by matching 1d index and data slices.
     For full documentation refer to :obj:`numpy.take_along_axis`.
@@ -582,32 +539,30 @@ def take_along_axis(arr, indices, axis):
     :obj:`put_along_axis` : Put values into the destination array by matching 1d index and data slices.
     """
 
-    if not use_origin_backend(arr):
-        if not isinstance(arr, dparray):
-            pass
-        elif not isinstance(indices, dparray):
-            pass
-        elif arr.ndim != indices.ndim:
+    x1_desc = dpnp.get_dpnp_descriptor(x1)
+    indices_desc = dpnp.get_dpnp_descriptor(indices)
+    if x1_desc and indices_desc:
+        if x1_desc.ndim != indices_desc.ndim:
             pass
         elif not isinstance(axis, int):
             pass
-        elif axis >= arr.ndim:
+        elif axis >= x1_desc.ndim:
             pass
-        elif arr.ndim == indices.ndim:
+        elif x1_desc.ndim == indices_desc.ndim:
             val_list = []
-            for i in list(indices.shape)[:-1]:
+            for i in list(indices_desc.shape)[:-1]:
                 if i == 1:
                     val_list.append(True)
                 else:
                     val_list.append(False)
             if not all(val_list):
                 pass
             else:
-                return dpnp_take_along_axis(arr, indices, axis)
+                return dpnp_take_along_axis(x1, indices, axis)
         else:
-            return dpnp_take_along_axis(arr, indices, axis)
+            return dpnp_take_along_axis(x1, indices, axis)
 
-    return call_origin(numpy.take_along_axis, arr, indices, axis)
+    return call_origin(numpy.take_along_axis, x1, indices, axis)
 
 
 def tril_indices(n, k=0, m=None):
@@ -644,7 +599,7 @@ def tril_indices(n, k=0, m=None):
     return call_origin(numpy.tril_indices, n, k, m)
 
 
-def tril_indices_from(arr, k=0):
+def tril_indices_from(x1, k=0):
     """
     Return the indices for the lower-triangle of arr.
     See `tril_indices` for full details.
@@ -659,13 +614,12 @@ def tril_indices_from(arr, k=0):
         Diagonal offset (see `tril` for details).
     """
 
-    is_arr_dparray = isinstance(arr, dparray)
-
-    if (not use_origin_backend(arr) and is_arr_dparray):
+    x1_desc = dpnp.get_dpnp_descriptor(x1)
+    if x1_desc:
         if isinstance(k, int):
-            return dpnp_tril_indices_from(arr, k)
+            return dpnp_tril_indices_from(x1_desc, k)
 
-    return call_origin(numpy.tril_indices_from, arr, k)
+    return call_origin(numpy.tril_indices_from, x1, k)
 
 
 def triu_indices(n, k=0, m=None):
@@ -702,7 +656,7 @@ def triu_indices(n, k=0, m=None):
     return call_origin(numpy.triu_indices, n, k, m)
 
 
-def triu_indices_from(arr, k=0):
+def triu_indices_from(x1, k=0):
     """
     Return the indices for the lower-triangle of arr.
     See `tril_indices` for full details.
@@ -717,10 +671,9 @@ def triu_indices_from(arr, k=0):
         Diagonal offset (see `tril` for details).
     """
 
-    is_arr_dparray = isinstance(arr, dparray)
-
-    if (not use_origin_backend(arr) and is_arr_dparray):
+    x1_desc = dpnp.get_dpnp_descriptor(x1)
+    if x1_desc:
         if isinstance(k, int):
-            return dpnp_triu_indices_from(arr, k)
+            return dpnp_triu_indices_from(x1_desc, k)
 
-    return call_origin(numpy.triu_indices_from, arr, k)
+    return call_origin(numpy.triu_indices_from, x1, k)
