Advanced indices don't broadcast if 1 array passed

- _mock removed from indexing methods

Author: ndgrigorian

dpctl/tensor/_copy_utils.py

@@ -430,7 +430,7 @@ def _mock_nonzero(ary):
     return tuple(dpt.asarray(i, usm_type=usm_type, sycl_queue=q) for i in nz)
 
 
-def _mock_take_multi_index(ary, inds, p):
+def _take_multi_index(ary, inds, p):
     if not isinstance(ary, dpt.usm_ndarray):
         raise TypeError
     queues_ = [
@@ -439,6 +439,8 @@ def _mock_take_multi_index(ary, inds, p):
     usm_types_ = [
         ary.usm_type,
     ]
+    if not isinstance(inds, list) and not isinstance(inds, tuple):
+        inds = (inds,)
     all_integers = True
     for ind in inds:
         queues_.append(ind.sycl_queue)
@@ -452,7 +454,8 @@ def _mock_take_multi_index(ary, inds, p):
         raise IndexError(
             "arrays used as indices must be of integer (or boolean) type"
         )
-    inds = dpt.broadcast_arrays(*inds)
+    if (len(inds) > 1):
+        inds = dpt.broadcast_arrays(*inds)
     ary_ndim = ary.ndim
     if ary_ndim > 0:
         p = operator.index(p)
@@ -505,7 +508,7 @@ def _mock_place(ary, ary_mask, p, vals):
     return
 
 
-def _mock_put_multi_index(ary, inds, p, vals):
+def _put_multi_index(ary, inds, p, vals):
     if isinstance(vals, dpt.usm_ndarray):
         queues_ = [ary.sycl_queue, vals.sycl_queue]
         usm_types_ = [ary.usm_type, vals.usm_type]
@@ -516,6 +519,8 @@ def _mock_put_multi_index(ary, inds, p, vals):
         usm_types_ = [
             ary.usm_type,
         ]
+    if not isinstance(inds, list) and not isinstance(inds, tuple):
+        inds = (inds,)
     all_integers = True
     for ind in inds:
         if not isinstance(ind, dpt.usm_ndarray):
@@ -536,8 +541,8 @@ def _mock_put_multi_index(ary, inds, p, vals):
         raise IndexError(
             "arrays used as indices must be of integer (or boolean) type"
         )
-
-    inds = dpt.broadcast_arrays(*inds)
+    if (len(inds) > 1):
+        inds = dpt.broadcast_arrays(*inds)
     ary_ndim = ary.ndim
     if ary_ndim > 0:
         p = operator.index(p)

dpctl/tensor/_indexing_functions.py

@@ -79,7 +79,8 @@ def take(x, indices, /, *, axis=None, mode="clip"):
             )
         axis = 0
 
-    indices = dpt.broadcast_arrays(*indices)
+    if len(indices) > 1:
+        indices = dpt.broadcast_arrays(*indices)
     if x_ndim > 0:
         axis = operator.index(axis)
         axis = normalize_axis_index(axis, x_ndim)
@@ -149,10 +150,13 @@ def put(x, indices, vals, /, *, axis=None, mode="clip"):
 
     # when axis is none, array is treated as 1D
     if axis is None:
-        x = dpt.reshape(x, (x.size,), copy=False)
-        axis = 0
-
-    indices = dpt.broadcast_arrays(*indices)
+        try:
+            x = dpt.reshape(x, (x.size,), copy=False)
+            axis = 0
+        except ValueError:
+            raise ValueError("Cannot create 1D view of array")
+    if len(indices) > 1:
+        indices = dpt.broadcast_arrays(*indices)
     x_ndim = x.ndim
     if x_ndim > 0:
         axis = operator.index(axis)

dpctl/tensor/_usmarray.pyx

@@ -673,7 +673,7 @@ cdef class usm_ndarray:
         from ._copy_utils import (
             _mock_extract,
             _mock_nonzero,
-            _mock_take_multi_index,
+            _take_multi_index,
         )
         if len(adv_ind) == 1 and adv_ind[0].dtype == dpt_bool:
             return _mock_extract(res, adv_ind[0], adv_ind_start_p)
@@ -685,9 +685,9 @@ cdef class usm_ndarray:
                     adv_ind_int.extend(_mock_nonzero(ind))
                 else:
                     adv_ind_int.append(ind)
-            return _mock_take_multi_index(res, tuple(adv_ind_int), adv_ind_start_p)
+            return _take_multi_index(res, tuple(adv_ind_int), adv_ind_start_p)
 
-        return _mock_take_multi_index(res, adv_ind, adv_ind_start_p)
+        return _take_multi_index(res, adv_ind, adv_ind_start_p)
 
 
     def to_device(self, target):
@@ -1021,7 +1021,7 @@ cdef class usm_ndarray:
             _copy_from_usm_ndarray_to_usm_ndarray,
             _mock_nonzero,
             _mock_place,
-            _mock_put_multi_index,
+            _put_multi_index,
         )
 
         adv_ind = _meta[3]
@@ -1064,10 +1064,10 @@ cdef class usm_ndarray:
                     adv_ind_int.extend(_mock_nonzero(ind))
                 else:
                     adv_ind_int.append(ind)
-            _mock_put_multi_index(Xv, tuple(adv_ind_int), adv_ind_start_p, rhs)
+            _put_multi_index(Xv, tuple(adv_ind_int), adv_ind_start_p, rhs)
             return
 
-        _mock_put_multi_index(Xv, adv_ind, adv_ind_start_p, rhs)
+        _put_multi_index(Xv, adv_ind, adv_ind_start_p, rhs)
         return