10360125231

Committed 12 Aug 2024 10:10PM UTC coverage: 87.751% (-0.2%) from 87.93%

Build # 10360125231

Build Type

Pull #1798

github

Committed by

web-flow

Commit Message

Merge 819a86500 into 43f5aeae1

Pull Request Pull Request #1798: Implements `dpctl.tensor.put_along_axis`

Run Details

3451 of 3973 branches covered (86.86%)

Branch coverage included in aggregate %.

10 of 32 new or added lines in 2 files covered. (31.25%)

1 existing line in 1 file now uncovered.

11464 of 13024 relevant lines covered (88.02%)

7115.01 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

88.09

/dpctl/tensor/_indexing_functions.py

#                       Data Parallel Control (dpctl)
#
#  Copyright 2020-2024 Intel 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 operator

import dpctl
import dpctl.tensor as dpt
import dpctl.tensor._tensor_impl as ti
import dpctl.utils

from ._copy_utils import (
    _extract_impl,
    _nonzero_impl,
    _put_multi_index,
    _take_multi_index,
)
from ._numpy_helper import normalize_axis_index


def _get_indexing_mode(name):
    modes = {"wrap": 0, "clip": 1}
    try:
        return modes[name]
    except KeyError:
        raise ValueError(
            "`mode` must be `wrap` or `clip`." "Got `{}`.".format(name)
        )


def take(x, indices, /, *, axis=None, mode="wrap"):
    """take(x, indices, axis=None, mode="wrap")

    Takes elements from an array along a given axis at given indices.

    Args:
        x (usm_ndarray):
            The array that elements will be taken from.
        indices (usm_ndarray):
            One-dimensional array of indices.
        axis (int, optional):
            The axis along which the values will be selected.
            If ``x`` is one-dimensional, this argument is optional.
            Default: ``None``.
        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:
          Array with shape
          ``x.shape[:axis] + indices.shape + x.shape[axis + 1:]``
          filled with elements from ``x``.
    """
    if not isinstance(x, dpt.usm_ndarray):
        raise TypeError(
            "Expected instance of `dpt.usm_ndarray`, got `{}`.".format(type(x))
        )

    if not isinstance(indices, dpt.usm_ndarray):
        raise TypeError(
            "`indices` expected `dpt.usm_ndarray`, got `{}`.".format(
                type(indices)
            )
        )
    if indices.dtype.kind not in "ui":
        raise IndexError(
            "`indices` expected integer data type, got `{}`".format(
                indices.dtype
            )
        )
    if indices.ndim != 1:
        raise ValueError(
            "`indices` expected a 1D array, got `{}`".format(indices.ndim)
        )
    exec_q = dpctl.utils.get_execution_queue([x.sycl_queue, indices.sycl_queue])
    if exec_q is None:
        raise dpctl.utils.ExecutionPlacementError
    res_usm_type = dpctl.utils.get_coerced_usm_type(
        [x.usm_type, indices.usm_type]
    )

    mode = _get_indexing_mode(mode)

    x_ndim = x.ndim
    if axis is None:
        if x_ndim > 1:
            raise ValueError(
                "`axis` cannot be `None` for array of dimension `{}`".format(
                    x_ndim
                )
            )
        axis = 0

    if x_ndim > 0:
        axis = normalize_axis_index(operator.index(axis), x_ndim)
        x_sh = x.shape
        if x_sh[axis] == 0 and indices.size != 0:
            raise IndexError("cannot take non-empty indices from an empty axis")
        res_shape = x.shape[:axis] + indices.shape + x.shape[axis + 1 :]
    else:
        if axis != 0:
            raise ValueError("`axis` must be 0 for an array of dimension 0.")
        res_shape = indices.shape

    res = dpt.empty(
        res_shape, dtype=x.dtype, usm_type=res_usm_type, sycl_queue=exec_q
    )

    _manager = dpctl.utils.SequentialOrderManager[exec_q]
    deps_ev = _manager.submitted_events
    hev, take_ev = ti._take(
        x, (indices,), res, axis, mode, sycl_queue=exec_q, depends=deps_ev
    )
    _manager.add_event_pair(hev, take_ev)

    return res


def put(x, indices, vals, /, *, axis=None, mode="wrap"):
    """put(x, indices, vals, axis=None, mode="wrap")

    Puts values into an array along a given axis at given indices.

    Args:
        x (usm_ndarray):
            The array the values will be put into.
        indices (usm_ndarray):
            One-dimensional array of indices.
        vals (usm_ndarray):
            Array of values to be put into ``x``.
            Must be broadcastable to the result shape
            ``x.shape[:axis] + indices.shape + x.shape[axis+1:]``.
        axis (int, optional):
            The axis along which the values will be placed.
            If ``x`` is one-dimensional, this argument is optional.
            Default: ``None``.
        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"``.

    .. note::

        If input array ``indices`` contains duplicates, a race condition
        occurs, and the value written into corresponding positions in ``x``
        may vary from run to run. Preserving sequential semantics in handing
        the duplicates to achieve deterministic behavior requires additional
        work, e.g.

        :Example:

            .. code-block:: python

                from dpctl import tensor as dpt

                def put_vec_duplicates(vec, ind, vals):
                    "Put values into vec, handling possible duplicates in ind"
                    assert vec.ndim, ind.ndim, vals.ndim == 1, 1, 1

                    # find positions of last occurences of each
                    # unique index
                    ind_flipped = dpt.flip(ind)
                    ind_uniq = dpt.unique_all(ind_flipped).indices
                    has_dups = len(ind) != len(ind_uniq)

                    if has_dups:
                        ind_uniq = dpt.subtract(vec.size - 1, ind_uniq)
                        ind = dpt.take(ind, ind_uniq)
                        vals = dpt.take(vals, ind_uniq)

                    dpt.put(vec, ind, vals)

                n = 512
                ind = dpt.concat((dpt.arange(n), dpt.arange(n, -1, step=-1)))
                x = dpt.zeros(ind.size, dtype="int32")
                vals = dpt.arange(ind.size, dtype=x.dtype)

                # Values corresponding to last positions of
                # duplicate indices are written into the vector x
                put_vec_duplicates(x, ind, vals)

                parts = (vals[-1:-n-2:-1], dpt.zeros(n, dtype=x.dtype))
                expected = dpt.concat(parts)
                assert dpt.all(x == expected)
    """
    if not isinstance(x, dpt.usm_ndarray):
        raise TypeError(
            "Expected instance of `dpt.usm_ndarray`, got `{}`.".format(type(x))
        )
    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)
        )
    if indices.dtype.kind not in "ui":
        raise IndexError(
            "`indices` expected integer data type, got `{}`".format(
                indices.dtype
            )
        )
    exec_q = dpctl.utils.get_execution_queue(queues_)
    if exec_q is None:
        raise dpctl.utils.ExecutionPlacementError
    vals_usm_type = dpctl.utils.get_coerced_usm_type(usm_types_)

    mode = _get_indexing_mode(mode)

    x_ndim = x.ndim
    if axis is None:
        if x_ndim > 1:
            raise ValueError(
                "`axis` cannot be `None` for array of dimension `{}`".format(
                    x_ndim
                )
            )
        axis = 0

    if x_ndim > 0:
        axis = normalize_axis_index(operator.index(axis), x_ndim)
        x_sh = x.shape
        if x_sh[axis] == 0 and indices.size != 0:
            raise IndexError("cannot take non-empty indices from an empty axis")
        val_shape = x.shape[:axis] + indices.shape + x.shape[axis + 1 :]
    else:
        if axis != 0:
            raise ValueError("`axis` must be 0 for an array of dimension 0.")
        val_shape = indices.shape

    if not isinstance(vals, dpt.usm_ndarray):
        vals = dpt.asarray(
            vals, dtype=x.dtype, usm_type=vals_usm_type, sycl_queue=exec_q
        )
    # choose to throw here for consistency with `place`
    if vals.size == 0:
        raise ValueError(
            "cannot put into non-empty indices along an empty axis"
        )
    if vals.dtype == x.dtype:
        rhs = vals
    else:
        rhs = dpt.astype(vals, x.dtype)
    rhs = dpt.broadcast_to(rhs, val_shape)

    _manager = dpctl.utils.SequentialOrderManager[exec_q]
    deps_ev = _manager.submitted_events
    hev, put_ev = ti._put(
        x, (indices,), rhs, axis, mode, sycl_queue=exec_q, depends=deps_ev
    )
    _manager.add_event_pair(hev, put_ev)


def extract(condition, arr):
    """extract(condition, arr)

    Returns the elements of an array that satisfies the condition.

    If ``condition`` is boolean ``dpctl.tensor.extract`` is
    equivalent to ``arr[condition]``.

    Note that ``dpctl.tensor.place`` does the opposite of
    ``dpctl.tensor.extract``.

    Args:
       conditions (usm_ndarray):
            An array whose non-zero or ``True`` entries indicate the element
            of ``arr`` to extract.

       arr (usm_ndarray):
            Input array of the same size as ``condition``.

    Returns:
        usm_ndarray:
            Rank 1 array of values from ``arr`` where ``condition`` is
            ``True``.
    """
    if not isinstance(condition, dpt.usm_ndarray):
        raise TypeError(
            "Expecting dpctl.tensor.usm_ndarray type, " f"got {type(condition)}"
        )
    if not isinstance(arr, dpt.usm_ndarray):
        raise TypeError(
            "Expecting dpctl.tensor.usm_ndarray type, " f"got {type(arr)}"
        )
    exec_q = dpctl.utils.get_execution_queue(
        (
            condition.sycl_queue,
            arr.sycl_queue,
        )
    )
    if exec_q is None:
        raise dpctl.utils.ExecutionPlacementError
    if condition.shape != arr.shape:
        raise ValueError("Arrays are not of the same size")
    return _extract_impl(arr, condition)


def place(arr, mask, vals):
    """place(arr, mask, vals)

    Change elements of an array based on conditional and input values.

    If ``mask`` is boolean ``dpctl.tensor.place`` is
    equivalent to ``arr[condition] = vals``.

    Args:
        arr (usm_ndarray):
            Array to put data into.
        mask (usm_ndarray):
            Boolean mask array. Must have the same size as ``arr``.
        vals (usm_ndarray, sequence):
            Values to put into ``arr``. Only the first N elements are
            used, where N is the number of True values in ``mask``. If
            ``vals`` is smaller than N, it will be repeated, and if
            elements of ``arr`` are to be masked, this sequence must be
            non-empty. Array ``vals`` must be one dimensional.
    """
    if not isinstance(arr, dpt.usm_ndarray):
        raise TypeError(
            "Expecting dpctl.tensor.usm_ndarray type, " f"got {type(arr)}"
        )
    if not isinstance(mask, dpt.usm_ndarray):
        raise TypeError(
            "Expecting dpctl.tensor.usm_ndarray type, " f"got {type(mask)}"
        )
    if not isinstance(vals, dpt.usm_ndarray):
        raise TypeError(
            "Expecting dpctl.tensor.usm_ndarray type, " f"got {type(vals)}"
        )
    exec_q = dpctl.utils.get_execution_queue(
        (
            arr.sycl_queue,
            mask.sycl_queue,
            vals.sycl_queue,
        )
    )
    if exec_q is None:
        raise dpctl.utils.ExecutionPlacementError
    if arr.shape != mask.shape or vals.ndim != 1:
        raise ValueError("Array sizes are not as required")
    cumsum = dpt.empty(mask.size, dtype="i8", sycl_queue=exec_q)
    _manager = dpctl.utils.SequentialOrderManager[exec_q]
    deps_ev = _manager.submitted_events
    nz_count = ti.mask_positions(
        mask, cumsum, sycl_queue=exec_q, depends=deps_ev
    )
    if nz_count == 0:
        return
    if vals.size == 0:
        raise ValueError("Cannot insert from an empty array!")
    if vals.dtype == arr.dtype:
        rhs = vals
    else:
        rhs = dpt.astype(vals, arr.dtype)
    hev, pl_ev = ti._place(
        dst=arr,
        cumsum=cumsum,
        axis_start=0,
        axis_end=mask.ndim,
        rhs=rhs,
        sycl_queue=exec_q,
    )
    _manager.add_event_pair(hev, pl_ev)


def nonzero(arr):
    """nonzero(arr)

    Return the indices of non-zero elements.

    Returns a tuple of usm_ndarrays, one for each dimension
    of ``arr``, containing the indices of the non-zero elements
    in that dimension. The values of ``arr`` are always tested in
    row-major, C-style order.

    Args:
        arr (usm_ndarray):
            Input array, which has non-zero array rank.

    Returns:
        Tuple[usm_ndarray, ...]:
            Indices of non-zero array elements.
    """
    if not isinstance(arr, dpt.usm_ndarray):
        raise TypeError(
            "Expecting dpctl.tensor.usm_ndarray type, " f"got {type(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(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)
    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)


def put_along_axis(x, indices, vals, /, *, axis=-1, mode="wrap"):
    """
    Puts elements into 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``.
        vals (usm_ndarray):
            Array of values to be put into ``x``.
            Must be broadcastable to the shape of ``indices``.
        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"``.

    .. note::

        If input array ``indices`` contains duplicates, a race condition
        occurs, and the value written into corresponding positions in ``x``
        may vary from run to run. Preserving sequential semantics in handing
        the duplicates to achieve deterministic behavior requires additional
        work.

        See :func:`dpctl.tensor.put` for an example.
    """
    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)

1	# Data Parallel Control (dpctl)
2	#
3	# Copyright 2020-2024 Intel Corporation
4	#
5	# Licensed under the Apache License, Version 2.0 (the "License");
6	# you may not use this file except in compliance with the License.
7	# You may obtain a copy of the License at
8	#
9	# http://www.apache.org/licenses/LICENSE-2.0
10	#
11	# Unless required by applicable law or agreed to in writing, software
12	# distributed under the License is distributed on an "AS IS" BASIS,
13	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14	# See the License for the specific language governing permissions and
15	# limitations under the License.
16
17	import operator	1✔
18
19	import dpctl	1✔
20	import dpctl.tensor as dpt	1✔
21	import dpctl.tensor._tensor_impl as ti	1✔
22	import dpctl.utils	1✔
23
24	from ._copy_utils import (	1✔
25	_extract_impl,
26	_nonzero_impl,
27	_put_multi_index,
28	_take_multi_index,
29	)
30	from ._numpy_helper import normalize_axis_index	1✔
31
32
33	def _get_indexing_mode(name):	1✔
34	modes = {"wrap": 0, "clip": 1}	1✔
35	try:	1✔
36	return modes[name]	1✔
37	except KeyError:	1✔
38	raise ValueError(	1✔
39	"`mode` must be `wrap` or `clip`." "Got `{}`.".format(name)
40	)
41
42
43	def take(x, indices, /, *, axis=None, mode="wrap"):	1✔
44	"""take(x, indices, axis=None, mode="wrap")
45
46	Takes elements from an array along a given axis at given indices.
47
48	Args:
49	x (usm_ndarray):
50	The array that elements will be taken from.
51	indices (usm_ndarray):
52	One-dimensional array of indices.
53	axis (int, optional):
54	The axis along which the values will be selected.
55	If ``x`` is one-dimensional, this argument is optional.
56	Default: ``None``.
57	mode (str, optional):
58	How out-of-bounds indices will be handled. Possible values
59	are:
60
61	- ``"wrap"``: clamps indices to (``-n <= i < n``), then wraps
62	negative indices.
63	- ``"clip"``: clips indices to (``0 <= i < n``).
64
65	Default: ``"wrap"``.
66
67	Returns:
68	usm_ndarray:
69	Array with shape
70	``x.shape[:axis] + indices.shape + x.shape[axis + 1:]``
71	filled with elements from ``x``.
72	"""
73	if not isinstance(x, dpt.usm_ndarray):	1✔
74	raise TypeError(	1✔
75	"Expected instance of `dpt.usm_ndarray`, got `{}`.".format(type(x))
76	)
77
78	if not isinstance(indices, dpt.usm_ndarray):	1✔
79	raise TypeError(	1✔
80	"`indices` expected `dpt.usm_ndarray`, got `{}`.".format(
81	type(indices)
82	)
83	)
84	if indices.dtype.kind not in "ui":	1✔
85	raise IndexError(	1✔
86	"`indices` expected integer data type, got `{}`".format(
87	indices.dtype
88	)
89	)
90	if indices.ndim != 1:	1✔
91	raise ValueError(	1✔
92	"`indices` expected a 1D array, got `{}`".format(indices.ndim)
93	)
94	exec_q = dpctl.utils.get_execution_queue([x.sycl_queue, indices.sycl_queue])	1✔
95	if exec_q is None:	1✔
96	raise dpctl.utils.ExecutionPlacementError	1✔
97	res_usm_type = dpctl.utils.get_coerced_usm_type(	1✔
98	[x.usm_type, indices.usm_type]
99	)
100
101	mode = _get_indexing_mode(mode)	1✔
102
103	x_ndim = x.ndim	1✔
104	if axis is None:	1✔
105	if x_ndim > 1:	1✔
106	raise ValueError(	1✔
107	"`axis` cannot be `None` for array of dimension `{}`".format(
108	x_ndim
109	)
110	)
111	axis = 0	1✔
112
113	if x_ndim > 0:	1✔
114	axis = normalize_axis_index(operator.index(axis), x_ndim)	1✔
115	x_sh = x.shape	1✔
116	if x_sh[axis] == 0 and indices.size != 0:	1✔
117	raise IndexError("cannot take non-empty indices from an empty axis")	1✔
118	res_shape = x.shape[:axis] + indices.shape + x.shape[axis + 1 :]	1✔
119	else:
120	if axis != 0:	1✔
121	raise ValueError("`axis` must be 0 for an array of dimension 0.")	1✔
122	res_shape = indices.shape	1✔
123
124	res = dpt.empty(	1✔
125	res_shape, dtype=x.dtype, usm_type=res_usm_type, sycl_queue=exec_q
126	)
127
128	_manager = dpctl.utils.SequentialOrderManager[exec_q]	1✔
129	deps_ev = _manager.submitted_events	1✔
130	hev, take_ev = ti._take(	1✔
131	x, (indices,), res, axis, mode, sycl_queue=exec_q, depends=deps_ev
132	)
133	_manager.add_event_pair(hev, take_ev)	1✔
134
135	return res	1✔
136
137
138	def put(x, indices, vals, /, *, axis=None, mode="wrap"):	1✔
139	"""put(x, indices, vals, axis=None, mode="wrap")
140
141	Puts values into an array along a given axis at given indices.
142
143	Args:
144	x (usm_ndarray):
145	The array the values will be put into.
146	indices (usm_ndarray):
147	One-dimensional array of indices.
148	vals (usm_ndarray):
149	Array of values to be put into ``x``.
150	Must be broadcastable to the result shape
151	``x.shape[:axis] + indices.shape + x.shape[axis+1:]``.
152	axis (int, optional):
153	The axis along which the values will be placed.
154	If ``x`` is one-dimensional, this argument is optional.
155	Default: ``None``.
156	mode (str, optional):
157	How out-of-bounds indices will be handled. Possible values
158	are:
159
160	- ``"wrap"``: clamps indices to (``-n <= i < n``), then wraps
161	negative indices.
162	- ``"clip"``: clips indices to (``0 <= i < n``).
163
164	Default: ``"wrap"``.
165
166	.. note::
167
168	If input array ``indices`` contains duplicates, a race condition
169	occurs, and the value written into corresponding positions in ``x``
170	may vary from run to run. Preserving sequential semantics in handing
171	the duplicates to achieve deterministic behavior requires additional
172	work, e.g.
173
174	:Example:
175
176	.. code-block:: python
177
178	from dpctl import tensor as dpt
179
180	def put_vec_duplicates(vec, ind, vals):
181	"Put values into vec, handling possible duplicates in ind"
182	assert vec.ndim, ind.ndim, vals.ndim == 1, 1, 1
183
184	# find positions of last occurences of each
185	# unique index
186	ind_flipped = dpt.flip(ind)
187	ind_uniq = dpt.unique_all(ind_flipped).indices
188	has_dups = len(ind) != len(ind_uniq)
189
190	if has_dups:
191	ind_uniq = dpt.subtract(vec.size - 1, ind_uniq)
192	ind = dpt.take(ind, ind_uniq)
193	vals = dpt.take(vals, ind_uniq)
194
195	dpt.put(vec, ind, vals)
196
197	n = 512
198	ind = dpt.concat((dpt.arange(n), dpt.arange(n, -1, step=-1)))
199	x = dpt.zeros(ind.size, dtype="int32")
200	vals = dpt.arange(ind.size, dtype=x.dtype)
201
202	# Values corresponding to last positions of
203	# duplicate indices are written into the vector x
204	put_vec_duplicates(x, ind, vals)
205
206	parts = (vals[-1:-n-2:-1], dpt.zeros(n, dtype=x.dtype))
207	expected = dpt.concat(parts)
208	assert dpt.all(x == expected)
209	"""
210	if not isinstance(x, dpt.usm_ndarray):	1✔
211	raise TypeError(	1✔
212	"Expected instance of `dpt.usm_ndarray`, got `{}`.".format(type(x))
213	)
214	if not isinstance(indices, dpt.usm_ndarray):	1✔
215	raise TypeError(	1✔
216	"`indices` expected `dpt.usm_ndarray`, got `{}`.".format(
217	type(indices)
218	)
219	)
220	if isinstance(vals, dpt.usm_ndarray):	1✔
221	queues_ = [x.sycl_queue, indices.sycl_queue, vals.sycl_queue]	1✔
222	usm_types_ = [x.usm_type, indices.usm_type, vals.usm_type]	1✔
223	else:
224	queues_ = [x.sycl_queue, indices.sycl_queue]	1✔
225	usm_types_ = [x.usm_type, indices.usm_type]	1✔
226	if indices.ndim != 1:	1✔
227	raise ValueError(	1✔
228	"`indices` expected a 1D array, got `{}`".format(indices.ndim)
229	)
230	if indices.dtype.kind not in "ui":	1✔
231	raise IndexError(	1✔
232	"`indices` expected integer data type, got `{}`".format(
233	indices.dtype
234	)
235	)
236	exec_q = dpctl.utils.get_execution_queue(queues_)	1✔
237	if exec_q is None:	1✔
238	raise dpctl.utils.ExecutionPlacementError	1✔
239	vals_usm_type = dpctl.utils.get_coerced_usm_type(usm_types_)	1✔
240
241	mode = _get_indexing_mode(mode)	1✔
242
243	x_ndim = x.ndim	1✔
244	if axis is None:	1✔
245	if x_ndim > 1:	1✔
246	raise ValueError(	1✔
247	"`axis` cannot be `None` for array of dimension `{}`".format(
248	x_ndim
249	)
250	)
251	axis = 0	1✔
252
253	if x_ndim > 0:	1✔
254	axis = normalize_axis_index(operator.index(axis), x_ndim)	1✔
255	x_sh = x.shape	1✔
256	if x_sh[axis] == 0 and indices.size != 0:	1✔
257	raise IndexError("cannot take non-empty indices from an empty axis")	1✔
258	val_shape = x.shape[:axis] + indices.shape + x.shape[axis + 1 :]	1✔
259	else:
260	if axis != 0:	1✔
261	raise ValueError("`axis` must be 0 for an array of dimension 0.")	1✔
262	val_shape = indices.shape	1✔
263
264	if not isinstance(vals, dpt.usm_ndarray):	1✔
265	vals = dpt.asarray(	1✔
266	vals, dtype=x.dtype, usm_type=vals_usm_type, sycl_queue=exec_q
267	)
268	# choose to throw here for consistency with `place`
269	if vals.size == 0:	1✔
270	raise ValueError(	1✔
271	"cannot put into non-empty indices along an empty axis"
272	)
273	if vals.dtype == x.dtype:	1✔
274	rhs = vals	1✔
275	else:
276	rhs = dpt.astype(vals, x.dtype)	1✔
277	rhs = dpt.broadcast_to(rhs, val_shape)	1✔
278
279	_manager = dpctl.utils.SequentialOrderManager[exec_q]	1✔
280	deps_ev = _manager.submitted_events	1✔
281	hev, put_ev = ti._put(	1✔
282	x, (indices,), rhs, axis, mode, sycl_queue=exec_q, depends=deps_ev
283	)
284	_manager.add_event_pair(hev, put_ev)	1✔
285
286
287	def extract(condition, arr):	1✔
288	"""extract(condition, arr)
289
290	Returns the elements of an array that satisfies the condition.
291
292	If ``condition`` is boolean ``dpctl.tensor.extract`` is
293	equivalent to ``arr[condition]``.
294
295	Note that ``dpctl.tensor.place`` does the opposite of
296	``dpctl.tensor.extract``.
297
298	Args:
299	conditions (usm_ndarray):
300	An array whose non-zero or ``True`` entries indicate the element
301	of ``arr`` to extract.
302
303	arr (usm_ndarray):
304	Input array of the same size as ``condition``.
305
306	Returns:
307	usm_ndarray:
308	Rank 1 array of values from ``arr`` where ``condition`` is
309	``True``.
310	"""
311	if not isinstance(condition, dpt.usm_ndarray):	1✔
312	raise TypeError(	1✔
313	"Expecting dpctl.tensor.usm_ndarray type, " f"got {type(condition)}"
314	)
315	if not isinstance(arr, dpt.usm_ndarray):	1✔
316	raise TypeError(	1✔
317	"Expecting dpctl.tensor.usm_ndarray type, " f"got {type(arr)}"
318	)
319	exec_q = dpctl.utils.get_execution_queue(	1✔
320	(
321	condition.sycl_queue,
322	arr.sycl_queue,
323	)
324	)
325	if exec_q is None:	1✔
326	raise dpctl.utils.ExecutionPlacementError	1✔
327	if condition.shape != arr.shape:	1✔
328	raise ValueError("Arrays are not of the same size")	1✔
329	return _extract_impl(arr, condition)	1✔
330
331
332	def place(arr, mask, vals):	1✔
333	"""place(arr, mask, vals)
334
335	Change elements of an array based on conditional and input values.
336
337	If ``mask`` is boolean ``dpctl.tensor.place`` is
338	equivalent to ``arr[condition] = vals``.
339
340	Args:
341	arr (usm_ndarray):
342	Array to put data into.
343	mask (usm_ndarray):
344	Boolean mask array. Must have the same size as ``arr``.
345	vals (usm_ndarray, sequence):
346	Values to put into ``arr``. Only the first N elements are
347	used, where N is the number of True values in ``mask``. If
348	``vals`` is smaller than N, it will be repeated, and if
349	elements of ``arr`` are to be masked, this sequence must be
350	non-empty. Array ``vals`` must be one dimensional.
351	"""
352	if not isinstance(arr, dpt.usm_ndarray):	1✔
353	raise TypeError(	1✔
354	"Expecting dpctl.tensor.usm_ndarray type, " f"got {type(arr)}"
355	)
356	if not isinstance(mask, dpt.usm_ndarray):	1✔
357	raise TypeError(	1✔
358	"Expecting dpctl.tensor.usm_ndarray type, " f"got {type(mask)}"
359	)
360	if not isinstance(vals, dpt.usm_ndarray):	1✔
361	raise TypeError(	1✔
362	"Expecting dpctl.tensor.usm_ndarray type, " f"got {type(vals)}"
363	)
364	exec_q = dpctl.utils.get_execution_queue(	1✔
365	(
366	arr.sycl_queue,
367	mask.sycl_queue,
368	vals.sycl_queue,
369	)
370	)
371	if exec_q is None:	1✔
372	raise dpctl.utils.ExecutionPlacementError	1✔
373	if arr.shape != mask.shape or vals.ndim != 1:	1✔
374	raise ValueError("Array sizes are not as required")	1✔
375	cumsum = dpt.empty(mask.size, dtype="i8", sycl_queue=exec_q)	1✔
376	_manager = dpctl.utils.SequentialOrderManager[exec_q]	1✔
377	deps_ev = _manager.submitted_events	1✔
378	nz_count = ti.mask_positions(	1✔
379	mask, cumsum, sycl_queue=exec_q, depends=deps_ev
380	)
381	if nz_count == 0:	1✔
382	return	1✔
383	if vals.size == 0:	1✔
384	raise ValueError("Cannot insert from an empty array!")	1✔
385	if vals.dtype == arr.dtype:	1✔
386	rhs = vals	1✔
387	else:
388	rhs = dpt.astype(vals, arr.dtype)	1✔
389	hev, pl_ev = ti._place(	1✔
390	dst=arr,
391	cumsum=cumsum,
392	axis_start=0,
393	axis_end=mask.ndim,
394	rhs=rhs,
395	sycl_queue=exec_q,
396	)
397	_manager.add_event_pair(hev, pl_ev)	1✔
398
399
400	def nonzero(arr):	1✔
401	"""nonzero(arr)
402
403	Return the indices of non-zero elements.
404
405	Returns a tuple of usm_ndarrays, one for each dimension
406	of ``arr``, containing the indices of the non-zero elements
407	in that dimension. The values of ``arr`` are always tested in
408	row-major, C-style order.
409
410	Args:
411	arr (usm_ndarray):
412	Input array, which has non-zero array rank.
413
414	Returns:
415	Tuple[usm_ndarray, ...]:
416	Indices of non-zero array elements.
417	"""
418	if not isinstance(arr, dpt.usm_ndarray):	1✔
419	raise TypeError(	1✔
420	"Expecting dpctl.tensor.usm_ndarray type, " f"got {type(arr)}"
421	)
422	if arr.ndim == 0:	1✔
423	raise ValueError("Array of positive rank is expected")	1✔
424	return _nonzero_impl(arr)	1✔
425
426
427	def _range(sh_i, i, nd, q, usm_t, dt):	1✔
428	ind = dpt.arange(sh_i, dtype=dt, usm_type=usm_t, sycl_queue=q)	1✔
429	ind.shape = tuple(sh_i if i == j else 1 for j in range(nd))	1✔
430	return ind	1✔
431
432
433	def take_along_axis(x, indices, /, *, axis=-1, mode="wrap"):	1✔
434	"""
435	Returns elements from an array at the one-dimensional indices specified
436	by ``indices`` along a provided ``axis``.
437
438	Args:
439	x (usm_ndarray):
440	input array. Must be compatible with ``indices``, except for the
441	axis (dimension) specified by ``axis``.
442	indices (usm_ndarray):
443	array indices. Must have the same rank (i.e., number of dimensions)
444	as ``x``.
445	axis: int
446	axis along which to select values. If ``axis`` is negative, the
447	function determines the axis along which to select values by
448	counting from the last dimension. Default: ``-1``.
449	mode (str, optional):
450	How out-of-bounds indices will be handled. Possible values
451	are:
452
453	- ``"wrap"``: clamps indices to (``-n <= i < n``), then wraps
454	negative indices.
455	- ``"clip"``: clips indices to (``0 <= i < n``).
456
457	Default: ``"wrap"``.
458
459	Returns:
460	usm_ndarray:
461	an array having the same data type as ``x``. The returned array has
462	the same rank (i.e., number of dimensions) as ``x`` and a shape
463	determined according to :ref:`broadcasting`, except for the axis
464	(dimension) specified by ``axis`` whose size must equal the size
465	of the corresponding axis (dimension) in ``indices``.
466
467	Note:
468	Treatment of the out-of-bound indices in ``indices`` array is controlled
469	by the value of ``mode`` keyword.
470	"""
471	if not isinstance(x, dpt.usm_ndarray):	1✔
472	raise TypeError(f"Expected dpctl.tensor.usm_ndarray, got {type(x)}")	1✔
473	if not isinstance(indices, dpt.usm_ndarray):	1✔
474	raise TypeError(	1✔
475	f"Expected dpctl.tensor.usm_ndarray, got {type(indices)}"
476	)
477	x_nd = x.ndim	1✔
478	if x_nd != indices.ndim:	1✔
479	raise ValueError(	1✔
480	"Number of dimensions in the first and the second "
481	"argument arrays must be equal"
482	)
483	pp = normalize_axis_index(operator.index(axis), x_nd)	1✔
484	out_usm_type = dpctl.utils.get_coerced_usm_type(	1✔
485	(x.usm_type, indices.usm_type)
486	)
487	exec_q = dpctl.utils.get_execution_queue((x.sycl_queue, indices.sycl_queue))	1✔
488	if exec_q is None:	1✔
489	raise dpctl.utils.ExecutionPlacementError(	1✔
490	"Execution placement can not be unambiguously inferred "
491	"from input arguments. "
492	)
493	mode_i = _get_indexing_mode(mode)	1✔
494	indexes_dt = ti.default_device_index_type(exec_q.sycl_device)	1✔
495	_ind = tuple(	1✔
496	(
497	indices
498	if i == pp
499	else _range(x.shape[i], i, x_nd, exec_q, out_usm_type, indexes_dt)
500	)
501	for i in range(x_nd)
502	)
503	return _take_multi_index(x, _ind, 0, mode=mode_i)	1✔
504
505
506	def put_along_axis(x, indices, vals, /, *, axis=-1, mode="wrap"):	1✔
507	"""
508	Puts elements into an array at the one-dimensional indices specified by
509	``indices`` along a provided ``axis``.
510
511	Args:
512	x (usm_ndarray):
513	input array. Must be compatible with ``indices``, except for the
514	axis (dimension) specified by ``axis``.
515	indices (usm_ndarray):
516	array indices. Must have the same rank (i.e., number of dimensions)
517	as ``x``.
518	vals (usm_ndarray):
519	Array of values to be put into ``x``.
520	Must be broadcastable to the shape of ``indices``.
521	axis: int
522	axis along which to select values. If ``axis`` is negative, the
523	function determines the axis along which to select values by
524	counting from the last dimension. Default: ``-1``.
525	mode (str, optional):
526	How out-of-bounds indices will be handled. Possible values
527	are:
528
529	- ``"wrap"``: clamps indices to (``-n <= i < n``), then wraps
530	negative indices.
531	- ``"clip"``: clips indices to (``0 <= i < n``).
532
533	Default: ``"wrap"``.
534
535	.. note::
536
537	If input array ``indices`` contains duplicates, a race condition
538	occurs, and the value written into corresponding positions in ``x``
539	may vary from run to run. Preserving sequential semantics in handing
540	the duplicates to achieve deterministic behavior requires additional
541	work.
542
543	See :func:`dpctl.tensor.put` for an example.
544	"""
NEW 545	if not isinstance(x, dpt.usm_ndarray):	×
NEW 546	raise TypeError(f"Expected dpctl.tensor.usm_ndarray, got {type(x)}")	×
NEW 547	if not isinstance(indices, dpt.usm_ndarray):	×
NEW 548	raise TypeError(	×
549	f"Expected dpctl.tensor.usm_ndarray, got {type(indices)}"
550	)
NEW 551	x_nd = x.ndim	×
NEW 552	if x_nd != indices.ndim:	×
NEW 553	raise ValueError(	×
554	"Number of dimensions in the first and the second "
555	"argument arrays must be equal"
556	)
NEW 557	pp = normalize_axis_index(operator.index(axis), x_nd)	×
NEW 558	if isinstance(vals, dpt.usm_ndarray):	×
NEW 559	queues_ = [x.sycl_queue, indices.sycl_queue, vals.sycl_queue]	×
NEW 560	usm_types_ = [x.usm_type, indices.usm_type, vals.usm_type]	×
561	else:
NEW 562	queues_ = [x.sycl_queue, indices.sycl_queue]	×
NEW 563	usm_types_ = [x.usm_type, indices.usm_type]	×
NEW 564	exec_q = dpctl.utils.get_execution_queue(queues_)	×
NEW 565	if exec_q is None:	×
NEW 566	raise dpctl.utils.ExecutionPlacementError(	×
567	"Execution placement can not be unambiguously inferred "
568	"from input arguments. "
569	)
NEW 570	out_usm_type = dpctl.utils.get_coerced_usm_type(usm_types_)	×
NEW 571	mode_i = _get_indexing_mode(mode)	×
NEW 572	indexes_dt = ti.default_device_index_type(exec_q.sycl_device)	×
NEW 573	_ind = tuple(	×
574	(
575	indices
576	if i == pp
577	else _range(x.shape[i], i, x_nd, exec_q, out_usm_type, indexes_dt)
578	)
579	for i in range(x_nd)
580	)
NEW UNCOV 581	return _put_multi_index(x, _ind, 0, vals, mode=mode_i)	×

IntelPython / dpctl / 10360125231

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous