10205281485

Committed 01 Aug 2024 08:26PM UTC coverage: 87.45% (-0.2%) from 87.611%

Build # 10205281485

Build Type

Pull #1778

github

Committed by

web-flow

Commit Message

Merge 26785a401 into f98fe159d

Pull Request Pull Request #1778: Implement take_along_axis function per Python Array API

Run Details

3330 of 3860 branches covered (86.27%)

Branch coverage included in aggregate %.

4 of 23 new or added lines in 2 files covered. (17.39%)

11 existing lines in 1 file now uncovered.

11268 of 12833 relevant lines covered (87.8%)

7184.89 hits per line

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

87.35

/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, _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 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 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
            )
        )
    queues_.append(indices.sycl_queue)
    usm_types_.append(indices.usm_type)
    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
    if not isinstance(indices, dpt.usm_ndarray):
        raise TypeError
    x_nd = x.ndim
    if x_nd != indices.ndim:
        raise ValueError
    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)

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

IntelPython / dpctl / 10205281485

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