Take and put tweaks

- take_multi_index and put_multi_index logic for 0D arrays removed, adjusted a test accordingly
- take, put, take_multi_index, and put_multi_index axis type check and normalization only reassigns axis once

Author: ndgrigorian

dpctl/tensor/_copy_utils.py

@@ -454,16 +454,12 @@ def _take_multi_index(ary, inds, p):
         raise IndexError(
             "arrays used as indices must be of integer (or boolean) type"
         )
-    if (len(inds) > 1):
+    if len(inds) > 1:
         inds = dpt.broadcast_arrays(*inds)
     ary_ndim = ary.ndim
-    if ary_ndim > 0:
-        p = operator.index(p)
-        p = normalize_axis_index(p, ary_ndim)
+    p = normalize_axis_index(operator.index(p), ary_ndim)
 
-        res_shape = ary.shape[:p] + inds[0].shape + ary.shape[p + len(inds) :]
-    else:
-        res_shape = inds[0].shape
+    res_shape = ary.shape[:p] + inds[0].shape + ary.shape[p + len(inds) :]
     res_usm_type = dpctl.utils.get_coerced_usm_type(usm_types_)
     res = dpt.empty(
         res_shape, dtype=ary.dtype, usm_type=res_usm_type, sycl_queue=exec_q
@@ -541,15 +537,12 @@ def _put_multi_index(ary, inds, p, vals):
         raise IndexError(
             "arrays used as indices must be of integer (or boolean) type"
         )
-    if (len(inds) > 1):
+    if len(inds) > 1:
         inds = dpt.broadcast_arrays(*inds)
     ary_ndim = ary.ndim
-    if ary_ndim > 0:
-        p = operator.index(p)
-        p = normalize_axis_index(p, ary_ndim)
-        vals_shape = ary.shape[:p] + inds[0].shape + ary.shape[p + len(inds) :]
-    else:
-        vals_shape = inds[0].shape
+
+    p = normalize_axis_index(operator.index(p), ary_ndim)
+    vals_shape = ary.shape[:p] + inds[0].shape + ary.shape[p + len(inds) :]
 
     vals_usm_type = dpctl.utils.get_coerced_usm_type(usm_types_)
     if not isinstance(vals, dpt.usm_ndarray):

dpctl/tensor/_indexing_functions.py

@@ -82,8 +82,7 @@ def take(x, indices, /, *, axis=None, mode="clip"):
     if len(indices) > 1:
         indices = dpt.broadcast_arrays(*indices)
     if x_ndim > 0:
-        axis = operator.index(axis)
-        axis = normalize_axis_index(axis, x_ndim)
+        axis = normalize_axis_index(operator.index(axis), x_ndim)
         res_shape = (
             x.shape[:axis] + indices[0].shape + x.shape[axis + len(indices) :]
         )
@@ -154,13 +153,12 @@ def put(x, indices, vals, /, *, axis=None, mode="clip"):
             x = dpt.reshape(x, (x.size,), copy=False)
             axis = 0
         except ValueError:
-            raise ValueError("Cannot create 1D view of array")
+            raise ValueError("Cannot create 1D view of input array")
     if len(indices) > 1:
         indices = dpt.broadcast_arrays(*indices)
     x_ndim = x.ndim
     if x_ndim > 0:
-        axis = operator.index(axis)
-        axis = normalize_axis_index(axis, x_ndim)
+        axis = normalize_axis_index(operator.index(axis), x_ndim)
 
         val_shape = (
             x.shape[:axis] + indices[0].shape + x.shape[axis + len(indices) :]

dpctl/tests/test_usm_ndarray_indexing.py

@@ -546,7 +546,8 @@ def test_put_0d_val(data_dt):
     assert_array_equal(np.asarray(2, dtype=data_dt), dpt.asnumpy(x[0]))
 
     x = dpt.asarray(5, dtype=data_dt, sycl_queue=q)
-    x[ind] = 2
+    val = 2
+    dpt.put(x, ind, val)
     assert_array_equal(np.asarray(2, dtype=data_dt), dpt.asnumpy(x))