Adding support for take_along_axis and put_along_axis

cunumeric/deferred.py

@@ -739,6 +739,21 @@ def set_item(self, key, rhs):
                     rhs = rhs._copy_store(rhs.base)
                 rhs_store = rhs.base
 
+            # the case when rhs is a scalar and indices array contains
+            # a single value
+            # TODO this logic should be removed when copy accepts Futures
+            if rhs_store.kind == Future:
+                rhs_tmp = DeferredArray(
+                    self.runtime,
+                    base=rhs_store,
+                    dtype=rhs.dtype,
+                )
+                rhs_tmp2 = self._convert_future_to_store(rhs_tmp)
+                rhs_store = rhs_tmp2.base
+
+            if index_array.base.kind == Future:
+                index_array = self._convert_future_to_store(index_array)
+
             copy = self.context.create_copy()
             copy.set_target_indirect_out_of_range(False)
 

cunumeric/module.py

@@ -22,6 +22,7 @@
 
 import numpy as np
 import opt_einsum as oe  # type: ignore [import]
+from numpy.core.multiarray import normalize_axis_index
 from numpy.core.numeric import (  # type: ignore [attr-defined]
     normalize_axis_tuple,
 )
@@ -2498,6 +2499,139 @@ def take(
     return a.take(indices=indices, axis=axis, out=out, mode=mode)
 
 
+def _fill_fancy_index_for_along_axis_routines(
+    a_shape: NdShape, axis: int, indices: ndarray
+) -> tuple[ndarray, ...]:
+
+    # the logic below is base on the cupy implementation of
+    # the *_along_axis routines
+    ndim = len(a_shape)
+    fancy_index = []
+    for i, n in enumerate(a_shape):
+        if i == axis:
+            fancy_index.append(indices)
+        else:
+            ind_shape = (1,) * i + (-1,) + (1,) * (ndim - i - 1)
+            fancy_index.append(arange(n).reshape(ind_shape))
+    return tuple(fancy_index)
+
+
+@add_boilerplate("a", "indices")
+def take_along_axis(a: ndarray, indices: ndarray, axis: int) -> ndarray:
+    """
+    Take values from the input array by matching 1d index and data slices.
+
+    This iterates over matching 1d slices oriented along the specified axis in
+    the index and data arrays, and uses the former to look up values in the
+    latter. These slices can be different lengths.
+
+    Functions returning an index along an axis, like
+    :func:`cunumeric.argsort` and :func:`cunumeric.argpartition`,
+    produce suitable indices for this function.
+
+    Parameters
+    ----------
+    arr : ndarray (Ni..., M, Nk...)
+        Source array
+    indices : ndarray (Ni..., J, Nk...)
+        Indices to take along each 1d slice of `arr`. This must match the
+        dimension of arr, but dimensions Ni and Nj only need to broadcast
+        against `arr`.
+    axis : 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
+        :func:`cunumeric.sort` and :func:`cunumeric.argsort`.
+
+    Returns
+    -------
+    out: ndarray (Ni..., J, Nk...)
+        The indexed result.
+
+    See Also
+    --------
+    numpy.take_along_axis
+
+    Availability
+    --------
+    Multiple GPUs, Multiple CPUs
+    """
+    if not np.issubdtype(indices.dtype, np.integer):
+        raise TypeError("`indices` must be an integer array")
+
+    if axis is None:
+        if indices.ndim != 1:
+            raise ValueError("indices must be 1D if axis=None")
+        a = a.ravel()
+        axis = 0
+    else:
+        axis = normalize_axis_index(axis, a.ndim)
+
+    ndim = a.ndim
+    if ndim != indices.ndim:
+        raise ValueError(
+            "`indices` and `a` must have the same number of dimensions"
+        )
+    return a[_fill_fancy_index_for_along_axis_routines(a.shape, axis, indices)]
+
+
+@add_boilerplate("a", "indices", "values")
+def put_along_axis(
+    a: ndarray, indices: ndarray, values: ndarray, axis: int
+) -> ndarray:
+    """
+    Put values into the destination array by matching 1d index and data slices.
+
+    This iterates over matching 1d slices oriented along the specified axis in
+    the index and data arrays, and uses the former to place values into the
+    latter. These slices can be different lengths.
+
+    Functions returning an index along an axis, like :func:`cunumeric.argsort`
+    and :func:`cunumeric.argpartition`, produce suitable indices for
+    this function.
+
+    Parameters
+    ----------
+    arr : ndarray (Ni..., M, Nk...)
+        Destination array.
+    indices : ndarray (Ni..., J, Nk...)
+        Indices to change along each 1d slice of `arr`. This must match the
+        dimension of arr, but dimensions in Ni and Nj may be 1 to broadcast
+        against `arr`.
+    values : array_like (Ni..., J, Nk...)
+        values to insert at those indices. Its shape and dimension are
+        broadcast to match that of `indices`.
+    axis : int
+        The axis to take 1d slices along. If axis is None, the destination
+        array is treated as if a flattened 1d view had been created of it.
+
+    See Also
+    --------
+    numpy.put_along_axis
+
+    Availability
+    --------
+    Multiple GPUs, Multiple CPUs
+
+    """
+    if not np.issubdtype(indices.dtype, np.integer):
+        raise IndexError("`indices` must be an integer array")
+
+    if axis is None:
+        if indices.ndim != 1:
+            raise ValueError("indices must be 1D if axis=None")
+        a = a.ravel()
+        axis = 0
+    else:
+        axis = normalize_axis_index(axis, a.ndim)
+
+    if a.ndim != indices.ndim:
+        raise ValueError(
+            "`indices` and `a` must have the same number of dimensions"
+        )
+    ind = _fill_fancy_index_for_along_axis_routines(a.shape, axis, indices)
+    a[ind] = values
+
+
 @add_boilerplate("a")
 def choose(
     a: ndarray,

docs/cunumeric/source/api/indexing.rst

@@ -27,3 +27,13 @@ Indexing-like operations
    diag
    diagonal
    take
+   take_along_axis
+
+
+Inserting data into arrays
+--------------------------
+
+.. autosummary::
+   :toctree: generated/
+
+   put_along_axis

tests/integration/test_put_along_axis.py

@@ -0,0 +1,75 @@
+# Copyright 2022 NVIDIA Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import numpy as np
+import pytest
+from test_tools.generators import mk_seq_array
+
+import cunumeric as num
+from legate.core import LEGATE_MAX_DIM
+
+
+def test_3d():
+
+    x = mk_seq_array(np, (256, 256, 100))
+    x_num = mk_seq_array(num, (256, 256, 100))
+
+    indices = mk_seq_array(np, (256, 256, 10)) % 100
+    indices_num = num.array(indices)
+
+    res = np.put_along_axis(x, indices, -10, -1)
+    res_num = num.put_along_axis(x_num, indices_num, -10, -1)
+    assert np.array_equal(res_num, res)
+
+
+def test_None_axis():
+    x = mk_seq_array(np, (256, 256, 100))
+    x_num = mk_seq_array(num, (256, 256, 100))
+
+    # testig the case when axis = None
+    indices = mk_seq_array(np, (256,))
+    indices_num = num.array(indices)
+
+    res = np.put_along_axis(x, indices, 99, None)
+    res_num = num.put_along_axis(x_num, indices_num, 99, None)
+    assert np.array_equal(res_num, res)
+
+
+N = 10
+
+
+@pytest.mark.parametrize("ndim", range(1, LEGATE_MAX_DIM + 1))
+def test_ndim(ndim):
+    shape = (N,) * ndim
+    np_arr = mk_seq_array(np, shape)
+    num_arr = mk_seq_array(num, shape)
+    shape_idx = (1,) * ndim
+    np_indices = mk_seq_array(np, shape_idx) % N
+    num_indices = mk_seq_array(num, shape_idx) % N
+    for axis in range(ndim):
+        res_np = np.put_along_axis(np_arr, np_indices, 8, axis=axis)
+        res_num = num.put_along_axis(num_arr, num_indices, 8, axis=axis)
+        assert np.array_equal(res_num, res_np)
+    np_indices = mk_seq_array(np, (3,))
+    num_indices = mk_seq_array(num, (3,))
+    res_np = np.put_along_axis(np_arr, np_indices, 11, None)
+    res_num = num.put_along_axis(num_arr, num_indices, 11, None)
+    assert np.array_equal(res_num, res_np)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(sys.argv))

tests/integration/test_take_along_axis.py

@@ -0,0 +1,75 @@
+# Copyright 2022 NVIDIA Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import numpy as np
+import pytest
+from test_tools.generators import mk_seq_array
+
+import cunumeric as num
+from legate.core import LEGATE_MAX_DIM
+
+
+def test_3d():
+
+    x = mk_seq_array(np, (256, 256, 100))
+    x_num = mk_seq_array(num, (256, 256, 100))
+
+    indices = mk_seq_array(np, (256, 256, 10)) % 100
+    indices_num = num.array(indices)
+
+    res = np.take_along_axis(x, indices, -1)
+    res_num = num.take_along_axis(x_num, indices_num, -1)
+    assert np.array_equal(res_num, res)
+
+
+def test_None_axis():
+    x = mk_seq_array(np, (256, 256, 100))
+    x_num = mk_seq_array(num, (256, 256, 100))
+
+    # testig the case when axis = None
+    indices = mk_seq_array(np, (256,))
+    indices_num = num.array(indices)
+
+    res = np.take_along_axis(x, indices, None)
+    res_num = num.take_along_axis(x_num, indices_num, None)
+    assert np.array_equal(res_num, res)
+
+
+N = 10
+
+
+@pytest.mark.parametrize("ndim", range(1, LEGATE_MAX_DIM + 1))
+def test_ndim(ndim):
+    shape = (N,) * ndim
+    np_arr = mk_seq_array(np, shape)
+    num_arr = mk_seq_array(num, shape)
+    shape_idx = (1,) * ndim
+    np_indices = mk_seq_array(np, shape_idx) % N
+    num_indices = mk_seq_array(num, shape_idx) % N
+    for axis in range(ndim):
+        res_np = np.take_along_axis(np_arr, np_indices, axis=axis)
+        res_num = num.take_along_axis(num_arr, num_indices, axis=axis)
+        assert np.array_equal(res_num, res_np)
+    np_indices = mk_seq_array(np, (3,))
+    num_indices = mk_seq_array(num, (3,))
+    res_np = np.take_along_axis(np_arr, np_indices, None)
+    res_num = num.take_along_axis(num_arr, num_indices, None)
+    assert np.array_equal(res_num, res_np)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(sys.argv))