DOC: Add documentation for take_along_axis

Abel Aoun · Abel Aoun · commit 105ef10f1898 · 2024-04-26T15:42:16.000+02:00
diff --git a/dask/array/chunk.py b/dask/array/chunk.py
@@ -433,16 +433,38 @@ def getitem(obj, index):
     return result
 
 
-def take_along_axis_chunk(arr, indices, offset, x_size, axis):
-    # Needed when idx is unsigned
+def take_along_axis_chunk(
+    arr: np.ndarray, indices: np.ndarray, offset: np.ndarray, arr_size: int, axis: int
+):
+    """Slice an ndarray according to ndarray indices along an axis.
+
+    Parameters
+    ----------
+    arr: np.ndarray, dtype=Any
+        The data array.
+    indices: np.ndarray, dtype=int64
+        The indices of interest.
+    offset: np.ndarray, shape=(1, ), dtype=int64
+        Index of the first element along axis of the current chunk of arr
+    arr_size: int
+        Total size of the arr da.Array along axis
+    axis: int
+        The axis along which the indices are from.
+
+    Returns
+    -------
+    out: np.ndarray
+        The indexed arr.
+    """
+    # Needed when indices is unsigned
     indices = indices.astype(np.int64)
     # Normalize negative indices
-    indices = np.where(indices < 0, indices + x_size, indices)
+    indices = np.where(indices < 0, indices + arr_size, indices)
     # A chunk of the offset dask Array is a numpy array with shape (1, ).
     # It indicates the index of the first element along axis of the current
-    # chunk of x.
+    # chunk of arr.
     indices = indices - offset
-    # Drop elements of idx that do not fall inside the current chunk of x
+    # Drop elements of idx that do not fall inside the current chunk of arr.
     idx_filter = (indices >= 0) & (indices < arr.shape[axis])
     indices[~idx_filter] = 0
     res = np.take_along_axis(arr, indices, axis=axis)
diff --git a/dask/array/slicing.py b/dask/array/slicing.py
@@ -7,6 +7,7 @@
 from itertools import product
 from numbers import Integral, Number
 from operator import itemgetter
+from typing import TYPE_CHECKING
 
 import numpy as np
 from tlz import concat, memoize, merge, pluck
@@ -17,6 +18,9 @@
 from dask.highlevelgraph import HighLevelGraph
 from dask.utils import cached_cumsum, is_arraylike
 
+if TYPE_CHECKING:
+    from dask.array import Array
+
 colon = slice(None, None, None)
 
 
@@ -2164,7 +2168,23 @@ def setitem(x, v, indices):
     return x
 
 
-def take_along_axis(arr, indices, axis):
+def take_along_axis(arr: Array, indices: Array, axis: int):
+    """Slice a dask ndarray according to dask ndarray of indices along an axis.
+
+    Parameters
+    ----------
+    arr: dask.array.Array, dtype=Any
+        Data array.
+    indices: dask.array.Array, dtype=int64
+        Indices of interest.
+    axis:int
+        The axis along which the indices are from.
+
+    Returns
+    -------
+    out: dask.array.Array
+        The indexed arr.
+    """
     from dask.array.core import Array, blockwise, from_array
 
     if axis < 0:
@@ -2179,29 +2199,30 @@ def take_along_axis(arr, indices, axis):
     # e.g. chunks=(..., (5, 3, 4), ...) -> offset=[0, 5, 8]
     offset = np.roll(np.cumsum(arr.chunks[axis]), 1)
     offset[0] = 0
-    offset = from_array(offset, chunks=1)
+    da_offset = from_array(offset, chunks=1)
     # Tamper with the declared chunks of offset to make blockwise align it with
-    # x[axis]
-    offset = Array(offset.dask, offset.name, (arr.chunks[axis],), offset.dtype)
+    # arr[axis]
+    da_offset = Array(
+        da_offset.dask, da_offset.name, (arr.chunks[axis],), da_offset.dtype
+    )
     # Define axis labels for blockwise
-    x_axes = tuple(range(arr.ndim))
+    arr_axes = tuple(range(arr.ndim))
     idx_label = (arr.ndim,)  # arbitrary unused
-    index_axes = x_axes[:axis] + idx_label + x_axes[axis + 1 :]
+    index_axes = arr_axes[:axis] + idx_label + arr_axes[axis + 1 :]
     offset_axes = (axis,)
-    p_axes = x_axes[: axis + 1] + idx_label + x_axes[axis + 1 :]
-    # Calculate the cartesian product of every chunk of x vs
-    # every chunk of index
+    p_axes = arr_axes[: axis + 1] + idx_label + arr_axes[axis + 1 :]
+    # Compute take_along_axis for each chunk
+    # TODO: Add meta argument for blockwise ?
     p = blockwise(
         take_along_axis_chunk,
         p_axes,
         arr,
-        x_axes,
+        arr_axes,
         indices,
         index_axes,
-        offset,
+        da_offset,
         offset_axes,
-        # align_arrays=False,
-        x_size=arr.shape[axis],
+        arr_size=arr.shape[axis],
         axis=axis,
         dtype=arr.dtype,
     )
diff --git a/docs/source/array-slicing.rst b/docs/source/array-slicing.rst
@@ -10,6 +10,8 @@ supports the following:
 *  Slicing one :class:`~dask.array.Array` with an :class:`~dask.array.Array` of bools: ``x[x > 0]``
 *  Slicing one :class:`~dask.array.Array` with a zero or one-dimensional :class:`~dask.array.Array`
    of ints: ``a[b.argtopk(5)]``
+*  Slicing one :class:`~dask.array.Array` with a multi-dimensional :class:`~dask.array.Array` of ints.
+   This can be done using ``dask.array.slicing.take_along_axis``.
 
 However, it does not currently support the following:
 
@@ -19,7 +21,6 @@ However, it does not currently support the following:
    issue. Also, users interested in this should take a look at
    :attr:`~dask.array.Array.vindex`.
 
-*  Slicing one :class:`~dask.array.Array` with a multi-dimensional :class:`~dask.array.Array` of ints
 
 .. _array.slicing.efficiency:
 
