Implement take_along_axis function per Python Array API

The function is planned for Python Array API 2024.12 specification.

Author: oleksandr-pavlyk

dpctl/tensor/__init__.py

@@ -60,7 +60,14 @@
 )
 from dpctl.tensor._device import Device
 from dpctl.tensor._dlpack import from_dlpack
-from dpctl.tensor._indexing_functions import extract, nonzero, place, put, take
+from dpctl.tensor._indexing_functions import (
+    extract,
+    nonzero,
+    place,
+    put,
+    take,
+    take_along_axis,
+)
 from dpctl.tensor._linear_algebra_functions import (
     matmul,
     matrix_transpose,
@@ -376,4 +383,5 @@
     "nextafter",
     "diff",
     "count_nonzero",
+    "take_along_axis",
 ]

dpctl/tensor/_indexing_functions.py

@@ -21,7 +21,7 @@
 import dpctl.tensor._tensor_impl as ti
 import dpctl.utils
 
-from ._copy_utils import _extract_impl, _nonzero_impl
+from ._copy_utils import _extract_impl, _nonzero_impl, _take_multi_index
 from ._numpy_helper import normalize_axis_index
 
 
@@ -423,3 +423,77 @@ def nonzero(arr):
     if arr.ndim == 0:
         raise ValueError("Array of positive rank is expected")
     return _nonzero_impl(arr)
+
+
+def _range(sh_i, i, nd, q, usm_t, dt):
+    ind = dpt.arange(sh_i, dtype=dt, usm_type=usm_t, sycl_queue=q)
+    ind.shape = tuple(sh_i if i == j else 1 for j in range(nd))
+    return ind
+
+
+def take_along_axis(x, indices, /, *, axis=-1, mode="wrap"):
+    """
+    Returns elements from an array at the one-dimensional indices specified
+    by ``indices`` along a provided ``axis``.
+
+    Args:
+        x (usm_ndarray):
+            input array. Must be compatible with ``indices``, except for the
+            axis (dimension) specified by ``axis``.
+        indices (usm_ndarray):
+            array indices. Must have the same rank (i.e., number of dimensions)
+            as ``x``.
+        axis: int
+            axis along which to select values. If ``axis`` is negative, the
+            function determines the axis along which to select values by
+            counting from the last dimension. Default: ``-1``.
+        mode (str, optional):
+            How out-of-bounds indices will be handled. Possible values
+            are:
+
+            - ``"wrap"``: clamps indices to (``-n <= i < n``), then wraps
+              negative indices.
+            - ``"clip"``: clips indices to (``0 <= i < n``).
+
+            Default: ``"wrap"``.
+
+    Returns:
+        usm_ndarray:
+            an array having the same data type as ``x``. The returned array has
+            the same rank (i.e., number of dimensions) as ``x`` and a shape
+            determined according to :ref:`broadcasting`, except for the axis
+            (dimension) specified by ``axis`` whose size must equal the size
+            of the corresponding axis (dimension) in ``indices``.
+
+    Note:
+        Treatment of the out-of-bound indices in ``indices`` array is controlled
+        by the value of ``mode`` keyword.
+    """
+    if not isinstance(x, dpt.usm_ndarray):
+        raise TypeError
+    if not isinstance(indices, dpt.usm_ndarray):
+        raise TypeError
+    x_nd = x.ndim
+    if x_nd != indices.ndim:
+        raise ValueError
+    pp = normalize_axis_index(operator.index(axis), x_nd)
+    out_usm_type = dpctl.utils.get_coerced_usm_type(
+        (x.usm_type, indices.usm_type)
+    )
+    exec_q = dpctl.utils.get_execution_queue((x.sycl_queue, indices.sycl_queue))
+    if exec_q is None:
+        raise dpctl.utils.ExecutionPlacementError(
+            "Execution placement can not be unambiguously inferred "
+            "from input arguments. "
+        )
+    mode_i = _get_indexing_mode(mode)
+    indexes_dt = ti.default_device_index_type(exec_q.sycl_device)
+    _ind = tuple(
+        (
+            indices
+            if i == pp
+            else _range(x.shape[i], i, x_nd, exec_q, out_usm_type, indexes_dt)
+        )
+        for i in range(x_nd)
+    )
+    return _take_multi_index(x, _ind, 0, mode=mode_i)