Implement dpnp.take_along_axis through dpctl.tensor

Author: antonwolfy

dpnp/dpnp_iface_indexing.py

@@ -1490,7 +1490,7 @@ def take(a, indices, /, *, axis=None, out=None, mode="wrap"):
     return dpnp.get_result_array(result, out)
 
 
-def take_along_axis(a, indices, axis):
+def take_along_axis(a, indices, axis, mode="wrap"):
     """
     Take values from the input array by matching 1d index and data slices.
 
@@ -1511,15 +1511,24 @@ def take_along_axis(a, indices, axis):
         Indices to take along each 1d slice of `a`. This must match the
         dimension of the input array, but dimensions ``Ni`` and ``Nj``
         only need to broadcast against `a`.
-    axis : int
+    axis : {None, int}
         The axis to take 1d slices along. If axis is ``None``, the input
         array is treated as if it had first been flattened to 1d,
         for consistency with :obj:`dpnp.sort` and :obj:`dpnp.argsort`.
+    mode : {"wrap", "clip"}, optional
+        Specifies 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
     -------
     out : dpnp.ndarray
-        The indexed result.
+        The indexed result of the same data type as `a`.
 
     See Also
     --------
@@ -1579,12 +1588,21 @@ def take_along_axis(a, indices, axis):
 
     """
 
-    dpnp.check_supported_arrays_type(a, indices)
-
     if axis is None:
-        a = a.ravel()
+        dpnp.check_supported_arrays_type(indices)
+        if indices.ndim != 1:
+            raise ValueError(
+                "when axis=None, `indices` must have a single dimension."
+            )
 
-    return a[_build_along_axis_index(a, indices, axis)]
+        a = dpnp.ravel(a)
+        axis = 0
+
+    usm_a = dpnp.get_usm_ndarray(a)
+    usm_ind = dpnp.get_usm_ndarray(indices)
+
+    usm_res = dpt.take_along_axis(usm_a, usm_ind, axis=axis, mode=mode)
+    return dpnp_array._create_from_usm_ndarray(usm_res)
 
 
 def tril_indices(