Skip to content

Commit

Permalink
Update dpnp.place implementation to get rid of limitations for inpu…
Browse files Browse the repository at this point in the history
…t arguments (#1912)

* Remove limitations from dpnp.take implementation

* Add more test to cover specail cases and increase code coverage

* Applied pre-commit hook

* Corrected test_over_index

* Update docsctrings with resolving typos

* Use dpnp.reshape() to change shape and create dpnp array from usm_ndarray result

* Remove limitations from dpnp.place implementation

* Update relating tests

* Roll back changed in dpnp.vander

* Remove data sync at the end of function

* Remove data sync from dpnp.get_result_array()

* Fix typo in docstring

* Corrected a link to dpnp.copyto() in description
  • Loading branch information
antonwolfy authored Jul 10, 2024
1 parent ce26cf0 commit e353d7d
Show file tree
Hide file tree
Showing 5 changed files with 158 additions and 167 deletions.
50 changes: 50 additions & 0 deletions dpnp/dpnp_iface.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,7 @@
"array_equal",
"asnumpy",
"astype",
"as_usm_ndarray",
"check_limitations",
"check_supported_arrays_type",
"convert_single_elem_array_to_scalar",
Expand Down Expand Up @@ -247,6 +248,55 @@ def astype(x1, dtype, order="K", casting="unsafe", copy=True, device=None):
return dpnp_array._create_from_usm_ndarray(array_obj)


def as_usm_ndarray(a, dtype=None, device=None, usm_type=None, sycl_queue=None):
"""
Return :class:`dpctl.tensor.usm_ndarray` from input object `a`.
Parameters
----------
a : {array_like, scalar}
Input array or scalar.
dtype : {None, dtype}, optional
The desired dtype for the result array if new array is creating. If not
given, a default dtype will be used that can represent the values (by
considering Promotion Type Rule and device capabilities when necessary).
Default: ``None``.
device : {None, string, SyclDevice, SyclQueue}, optional
An array API concept of device where the result array is created if
required.
The `device` can be ``None`` (the default), an OneAPI filter selector
string, an instance of :class:`dpctl.SyclDevice` corresponding to
a non-partitioned SYCL device, an instance of :class:`dpctl.SyclQueue`,
or a `Device` object returned by
:obj:`dpnp.dpnp_array.dpnp_array.device` property.
Default: ``None``.
usm_type : {None, "device", "shared", "host"}, optional
The type of SYCL USM allocation for the result array if new array
is created.
Default: ``None``.
sycl_queue : {None, SyclQueue}, optional
A SYCL queue to use for result array allocation if required.
Default: ``None``.
Returns
-------
out : usm_ndarray
A dpctl USM ndarray from input array or scalar `a`.
If `a` is instance of :class:`dpnp.ndarray`
or :class:`dpctl.tensor.usm_ndarray`, no array allocation will be done
and `dtype`, `device`, `usm_type`, `sycl_queue` keywords
will be ignored.
"""

if is_supported_array_type(a):
return get_usm_ndarray(a)

return dpt.asarray(
a, dtype=dtype, device=device, usm_type=usm_type, sycl_queue=sycl_queue
)


def check_limitations(
order=None, subok=False, like=None, initial=None, where=True
):
Expand Down
8 changes: 3 additions & 5 deletions dpnp/dpnp_iface_histograms.py
Original file line number Diff line number Diff line change
Expand Up @@ -164,11 +164,9 @@ def _get_bin_edges(a, bins, range, usm_type):
"a and bins must be allocated on the same SYCL queue"
)

bin_edges = dpnp.get_usm_ndarray(bins)
else:
bin_edges = dpt.asarray(
bins, sycl_queue=sycl_queue, usm_type=usm_type
)
bin_edges = dpnp.as_usm_ndarray(
bins, usm_type=usm_type, sycl_queue=sycl_queue
)

if dpnp.any(bin_edges[:-1] > bin_edges[1:]):
raise ValueError(
Expand Down
90 changes: 67 additions & 23 deletions dpnp/dpnp_iface_indexing.py
Original file line number Diff line number Diff line change
Expand Up @@ -551,12 +551,12 @@ def extract(condition, a):
"""

usm_a = dpnp.get_usm_ndarray(a)
if not dpnp.is_supported_array_type(condition):
usm_cond = dpt.asarray(
condition, usm_type=a.usm_type, sycl_queue=a.sycl_queue
)
else:
usm_cond = dpnp.get_usm_ndarray(condition)
usm_cond = dpnp.as_usm_ndarray(
condition,
dtype=dpnp.bool,
usm_type=usm_a.usm_type,
sycl_queue=usm_a.sycl_queue,
)

if usm_cond.size != usm_a.size:
usm_a = dpt.reshape(usm_a, -1)
Expand Down Expand Up @@ -1011,30 +1011,74 @@ def nonzero(a):
)


def place(x, mask, vals, /):
def place(a, mask, vals):
"""
Change elements of an array based on conditional and input values.
Similar to ``dpnp.copyto(a, vals, where=mask)``, the difference is that
:obj:`dpnp.place` uses the first N elements of `vals`, where N is
the number of ``True`` values in `mask`, while :obj:`dpnp.copyto` uses
the elements where `mask` is ``True``.
Note that :obj:`dpnp.extract` does the exact opposite of :obj:`dpnp.place`.
For full documentation refer to :obj:`numpy.place`.
Limitations
-----------
Parameters `x`, `mask` and `vals` are supported either as
:class:`dpnp.ndarray` or :class:`dpctl.tensor.usm_ndarray`.
Otherwise the function will be executed sequentially on CPU.
Parameters
----------
a : {dpnp.ndarray, usm_ndarray}
Array to put data into.
mask : {array_like, scalar}
Boolean mask array. Must have the same size as `a`.
vals : {array_like, scalar}
Values to put into `a`. 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 `a` are to be masked, this
sequence must be non-empty.
See Also
--------
:obj:`dpnp.copyto` : Copies values from one array to another.
:obj:`dpnp.put` : Replaces specified elements of an array with given values.
:obj:`dpnp.take` : Take elements from an array along an axis.
:obj:`dpnp.extract` : Return the elements of an array that satisfy some
condition.
Examples
--------
>>> import dpnp as np
>>> a = np.arange(6).reshape(2, 3)
>>> np.place(a, a > 2, [44, 55])
>>> a
array([[ 0, 1, 2],
[44, 55, 44]])
"""

if (
dpnp.is_supported_array_type(x)
and dpnp.is_supported_array_type(mask)
and dpnp.is_supported_array_type(vals)
):
dpt_array = x.get_array() if isinstance(x, dpnp_array) else x
dpt_mask = mask.get_array() if isinstance(mask, dpnp_array) else mask
dpt_vals = vals.get_array() if isinstance(vals, dpnp_array) else vals
return dpt.place(dpt_array, dpt_mask, dpt_vals)

return call_origin(numpy.place, x, mask, vals, dpnp_inplace=True)
usm_a = dpnp.get_usm_ndarray(a)
usm_mask = dpnp.as_usm_ndarray(
mask,
dtype=dpnp.bool,
usm_type=usm_a.usm_type,
sycl_queue=usm_a.sycl_queue,
)
usm_vals = dpnp.as_usm_ndarray(
vals,
dtype=usm_a.dtype,
usm_type=usm_a.usm_type,
sycl_queue=usm_a.sycl_queue,
)

if usm_vals.ndim != 1:
# dpt.place supports only 1-D array of values
usm_vals = dpt.reshape(usm_vals, -1)

if usm_vals.dtype != usm_a.dtype:
# dpt.place casts values to a.dtype with "unsafe" rule,
# while numpy.place does that with "safe" casting rule
usm_vals = dpt.astype(usm_vals, usm_a.dtype, casting="safe", copy=False)

dpt.place(usm_a, usm_mask, usm_vals)


def put(a, ind, v, /, *, axis=None, mode="wrap"):
Expand Down
Loading

0 comments on commit e353d7d

Please sign in to comment.