Implements put_along_axis

Also makes minor tweaks to `take_along_axis`

Author: ndgrigorian

dpctl/tensor/__init__.py

@@ -65,6 +65,7 @@
     nonzero,
     place,
     put,
+    put_along_axis,
     take,
     take_along_axis,
 )
@@ -384,4 +385,5 @@
     "diff",
     "count_nonzero",
     "take_along_axis",
+    "put_along_axis",
 ]

dpctl/tensor/_copy_utils.py

@@ -938,13 +938,18 @@ def _place_impl(ary, ary_mask, vals, axis=0):
     return
 
 
-def _put_multi_index(ary, inds, p, vals):
+def _put_multi_index(ary, inds, p, vals, mode=0):
     if not isinstance(ary, dpt.usm_ndarray):
         raise TypeError(
             f"Expecting type dpctl.tensor.usm_ndarray, got {type(ary)}"
         )
     ary_nd = ary.ndim
     p = normalize_axis_index(operator.index(p), ary_nd)
+    mode = operator.index(mode)
+    if mode not in [0, 1]:
+        raise ValueError(
+            "Invalid value for mode keyword, only 0 or 1 is supported"
+        )
     if isinstance(vals, dpt.usm_ndarray):
         queues_ = [ary.sycl_queue, vals.sycl_queue]
         usm_types_ = [ary.usm_type, vals.usm_type]
@@ -1018,7 +1023,7 @@ def _put_multi_index(ary, inds, p, vals):
         ind=inds,
         val=rhs,
         axis_start=p,
-        mode=0,
+        mode=mode,
         sycl_queue=exec_q,
         depends=dep_ev,
     )

dpctl/tensor/_indexing_functions.py

@@ -21,7 +21,12 @@
 import dpctl.tensor._tensor_impl as ti
 import dpctl.utils
 
-from ._copy_utils import _extract_impl, _nonzero_impl, _take_multi_index
+from ._copy_utils import (
+    _extract_impl,
+    _nonzero_impl,
+    _put_multi_index,
+    _take_multi_index,
+)
 from ._numpy_helper import normalize_axis_index
 
 
@@ -206,22 +211,18 @@ def put_vec_duplicates(vec, ind, vals):
         raise TypeError(
             "Expected instance of `dpt.usm_ndarray`, got `{}`.".format(type(x))
         )
-    if isinstance(vals, dpt.usm_ndarray):
-        queues_ = [x.sycl_queue, vals.sycl_queue]
-        usm_types_ = [x.usm_type, vals.usm_type]
-    else:
-        queues_ = [
-            x.sycl_queue,
-        ]
-        usm_types_ = [
-            x.usm_type,
-        ]
     if not isinstance(indices, dpt.usm_ndarray):
         raise TypeError(
             "`indices` expected `dpt.usm_ndarray`, got `{}`.".format(
                 type(indices)
             )
         )
+    if isinstance(vals, dpt.usm_ndarray):
+        queues_ = [x.sycl_queue, indices.sycl_queue, vals.sycl_queue]
+        usm_types_ = [x.usm_type, indices.usm_type, vals.usm_type]
+    else:
+        queues_ = [x.sycl_queue, indices.sycl_queue]
+        usm_types_ = [x.usm_type, indices.usm_type]
     if indices.ndim != 1:
         raise ValueError(
             "`indices` expected a 1D array, got `{}`".format(indices.ndim)
@@ -232,7 +233,6 @@ def put_vec_duplicates(vec, ind, vals):
                 indices.dtype
             )
         )
-    queues_.append(indices.sycl_queue)
     usm_types_.append(indices.usm_type)
     exec_q = dpctl.utils.get_execution_queue(queues_)
     if exec_q is None:
@@ -491,8 +491,12 @@ def take_along_axis(x, indices, /, *, axis=-1, mode="wrap"):
             "Execution placement can not be unambiguously inferred "
             "from input arguments. "
         )
+    indexes_dt = indices.dtype
+    if indexes_dt.kind not in "ui":
+        raise IndexError(
+            "`indices` expected integer data type, got `{}`".format(indexes_dt)
+        )
     mode_i = _get_indexing_mode(mode)
-    indexes_dt = ti.default_device_index_type(exec_q.sycl_device)
     _ind = tuple(
         (
             indices
@@ -502,3 +506,80 @@ def take_along_axis(x, indices, /, *, axis=-1, mode="wrap"):
         for i in range(x_nd)
     )
     return _take_multi_index(x, _ind, 0, mode=mode_i)
+
+
+def put_along_axis(x, indices, vals, /, *, 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(f"Expected dpctl.tensor.usm_ndarray, got {type(x)}")
+    if not isinstance(indices, dpt.usm_ndarray):
+        raise TypeError(
+            f"Expected dpctl.tensor.usm_ndarray, got {type(indices)}"
+        )
+    x_nd = x.ndim
+    if x_nd != indices.ndim:
+        raise ValueError(
+            "Number of dimensions in the first and the second "
+            "argument arrays must be equal"
+        )
+    pp = normalize_axis_index(operator.index(axis), x_nd)
+    if isinstance(vals, dpt.usm_ndarray):
+        queues_ = [x.sycl_queue, indices.sycl_queue, vals.sycl_queue]
+        usm_types_ = [x.usm_type, indices.usm_type, vals.usm_type]
+    else:
+        queues_ = [x.sycl_queue, indices.sycl_queue]
+        usm_types_ = [x.usm_type, indices.usm_type]
+    exec_q = dpctl.utils.get_execution_queue(queues_)
+    if exec_q is None:
+        raise dpctl.utils.ExecutionPlacementError(
+            "Execution placement can not be unambiguously inferred "
+            "from input arguments. "
+        )
+    out_usm_type = dpctl.utils.get_coerced_usm_type(usm_types_)
+    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 _put_multi_index(x, _ind, 0, vals, mode=mode_i)