Merge pull request #2935 from IntelPython/technical-debt-tensor-indexing

Clean up technical debt in integer indexing Python bindings

Author: ndgrigorian

CHANGELOG.md

@@ -21,6 +21,7 @@ This release is compatible with NumPy 2.4.5.
 * Updated tests to align with NumPy 2.4.5 compatibility [gh-2920](https://github.com/IntelPython/dpnp/pull/2920)
 * Replaced `.pxi` includes in `dpnp.tensor` with modular `.pxd`/`.pyx` Cython imports [#2913](https://github.com/IntelPython/dpnp/pull/2913)
 * Reimplemented `dpnp.eye` and `dpnp.tensor.eye` with a branchless kernel [gh-2937](https://github.com/IntelPython/dpnp/pull/2937)
+* Cleaned up Python bindings for indexing functions, renaming `usm_ndarray_take` and `usm_ndarray_put` to `py_take` and `py_put` and refactoring validation [gh-2935](https://github.com/IntelPython/dpnp/pull/2935)
 
 ### Deprecated
 

dpnp/dpnp_iface_indexing.py

@@ -312,6 +312,7 @@ def _take_index(x, inds, axis, q, usm_type, out=None, mode=0):
         ind=(inds,),
         dst=out,
         axis_start=axis,
+        axis_end=axis_end,
         mode=mode,
         sycl_queue=q,
         depends=dep_evs,

dpnp/tensor/_copy_utils.py

@@ -475,6 +475,7 @@ def _put_multi_index(ary, inds, p, vals, mode=0):
         ind=inds,
         val=rhs,
         axis_start=p,
+        axis_end=p_end,
         mode=mode,
         sycl_queue=exec_q,
         depends=dep_ev,
@@ -527,6 +528,7 @@ def _take_multi_index(ary, inds, p, mode=0):
         ind=inds,
         dst=res,
         axis_start=p,
+        axis_end=p_end,
         mode=mode,
         sycl_queue=exec_q,
         depends=dep_ev,

dpnp/tensor/_indexing_functions.py

@@ -342,7 +342,14 @@ def put_vec_duplicates(vec, ind, vals):
     _manager = SequentialOrderManager[exec_q]
     deps_ev = _manager.submitted_events
     hev, put_ev = ti._put(
-        x, (indices,), rhs, axis, mode, sycl_queue=exec_q, depends=deps_ev
+        x,
+        (indices,),
+        rhs,
+        axis,
+        axis + 1,
+        mode,
+        sycl_queue=exec_q,
+        depends=deps_ev,
     )
     _manager.add_event_pair(hev, put_ev)
 
@@ -543,7 +550,14 @@ def take(x, indices, /, *, axis=None, out=None, mode="wrap"):
     _manager = SequentialOrderManager[exec_q]
     deps_ev = _manager.submitted_events
     hev, take_ev = ti._take(
-        x, (indices,), out, axis, mode, sycl_queue=exec_q, depends=deps_ev
+        x,
+        (indices,),
+        out,
+        axis,
+        axis + 1,
+        mode,
+        sycl_queue=exec_q,
+        depends=deps_ev,
     )
     _manager.add_event_pair(hev, take_ev)
 

dpnp/tensor/_searchsorted.py

@@ -160,6 +160,7 @@ def searchsorted(
             ind,
             res,
             axis,
+            axis + 1,
             wrap_out_of_bound_indices_mode,
             sycl_queue=q,
             depends=dep_evs,

dpnp/tensor/_set_functions.py

@@ -383,6 +383,7 @@ def unique_inverse(x):
         ind=(sorting_ids,),
         dst=s,
         axis_start=0,
+        axis_end=1,
         mode=0,
         sycl_queue=exec_q,
         depends=[sort_ev],
@@ -558,6 +559,7 @@ def unique_all(x: dpt.usm_ndarray) -> UniqueAllResult:
         ind=(sorting_ids,),
         dst=s,
         axis_start=0,
+        axis_end=1,
         mode=0,
         sycl_queue=exec_q,
         depends=[sort_ev],