Skip to content

Addresses gh-1360 with documentation #1653

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Apr 23, 2024
Merged

Addresses gh-1360 with documentation #1653

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
124458f
Addresses gh-1360 with documentation
oleksandr-pavlyk Apr 23, 2024
18df92a
Removed obsolete notice in description of indices
oleksandr-pavlyk Apr 23, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
156 changes: 102 additions & 54 deletions dpctl/tensor/_indexing_functions.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,25 +41,29 @@ def 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:
The axis along which the values will be selected.
If ``x`` is one-dimensional, this argument is optional.
Default: ``None``.
mode:
How out-of-bounds indices will be handled.
``"wrap"`` - clamps indices to (-n <= i < n), then wraps
negative indices.
``"clip"`` - clips indices to (0 <= i < n)
Default: ``"wrap"``.
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.
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(
Expand Down Expand Up @@ -128,30 +132,71 @@ def 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.

Note that if indices are not unique, a race
condition will result, and the value written to
``x`` will not be deterministic.
:py:func:`dpctl.tensor.unique` can be used to
guarantee unique elements in ``indices``.
vals:
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:
The axis along which the values will be placed.
If ``x`` is one-dimensional, this argument is optional.
Default: ``None``.
mode:
How out-of-bounds indices will be handled.
``"wrap"`` - clamps indices to (-n <= i < n), then wraps
negative indices.
``"clip"`` - clips indices to (0 <= i < n)
Default: ``"wrap"``.
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(
Expand Down Expand Up @@ -237,22 +282,24 @@ def extract(condition, arr):

Returns the elements of an array that satisfies the condition.

If `condition` is boolean ``dpctl.tensor.extract`` is
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.
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`.
Input array of the same size as ``condition``.

Returns:
usm_ndarray:
Rank 1 array of values from `arr` where `condition` is True.
Rank 1 array of values from ``arr`` where ``condition`` is
``True``.
"""
if not isinstance(condition, dpt.usm_ndarray):
raise TypeError(
Expand Down Expand Up @@ -280,20 +327,20 @@ def place(arr, mask, vals):

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

If `mask` is boolean ``dpctl.tensor.place`` is
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`.
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.
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(
Expand Down Expand Up @@ -345,13 +392,14 @@ def 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
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.
Expand Down