Skip to content
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.

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

Added device keyword argument to astype function #1870

Merged
merged 4 commits into from
Jun 14, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
21 changes: 19 additions & 2 deletions dpnp/dpnp_array.py
Original file line number Diff line number Diff line change
Expand Up @@ -562,7 +562,15 @@ def asnumpy(self):

return dpt.asnumpy(self._array_obj)

def astype(self, dtype, order="K", casting="unsafe", subok=True, copy=True):
def astype(
self,
dtype,
order="K",
casting="unsafe",
subok=True,
copy=True,
device=None,
):
"""
Copy the array with data type casting.

Expand Down Expand Up @@ -597,6 +605,13 @@ def astype(self, dtype, order="K", casting="unsafe", subok=True, copy=True):
this is set to ``False``, and the `dtype`, `order`, and `subok`
requirements are satisfied, the input array is returned instead of
a copy.
device : {None, string, SyclDevice, SyclQueue}, optional
An array API concept of device where the output array is created.
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``.

Returns
-------
Expand Down Expand Up @@ -626,7 +641,9 @@ def astype(self, dtype, order="K", casting="unsafe", subok=True, copy=True):
f"subok={subok} is currently not supported"
)

return dpnp.astype(self, dtype, order=order, casting=casting, copy=copy)
return dpnp.astype(
self, dtype, order=order, casting=casting, copy=copy, device=device
)

# 'base',
# 'byteswap',
Expand Down
11 changes: 9 additions & 2 deletions dpnp/dpnp_iface.py
Original file line number Diff line number Diff line change
Expand Up @@ -180,7 +180,7 @@ def asnumpy(a, order="C"):


# pylint: disable=redefined-outer-name
def astype(x1, dtype, order="K", casting="unsafe", copy=True):
def astype(x1, dtype, order="K", casting="unsafe", copy=True, device=None):
"""
Copy the array with data type casting.

Expand Down Expand Up @@ -213,6 +213,13 @@ def astype(x1, dtype, order="K", casting="unsafe", copy=True):
By default, ``astype`` always returns a newly allocated array. If this
is set to ``False``, and the `dtype`, `order`, and `subok` requirements
are satisfied, the input array is returned instead of a copy.
device : {None, string, SyclDevice, SyclQueue}, optional
An array API concept of device where the output array is created.
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``.

Returns
-------
Expand All @@ -228,7 +235,7 @@ def astype(x1, dtype, order="K", casting="unsafe", copy=True):

x1_obj = dpnp.get_usm_ndarray(x1)
array_obj = dpt.astype(
x1_obj, dtype, order=order, casting=casting, copy=copy
x1_obj, dtype, order=order, casting=casting, copy=copy, device=device
)

# return x1 if dpctl returns a zero copy of x1_obj
Expand Down
19 changes: 19 additions & 0 deletions tests/test_sycl_queue.py
Original file line number Diff line number Diff line change
Expand Up @@ -2211,3 +2211,22 @@ def test_histogram_bin_edges(weights, device):

edges_queue = result_edges.sycl_queue
assert_sycl_queue_equal(edges_queue, iv.sycl_queue)


@pytest.mark.parametrize(
"device_x",
valid_devices,
ids=[device.filter_string for device in valid_devices],
)
@pytest.mark.parametrize(
"device_y",
valid_devices,
ids=[device.filter_string for device in valid_devices],
)
def test_astype(device_x, device_y):
x = dpnp.array([1, 2, 3], dtype="i4", device=device_x)
y = dpnp.astype(x, dtype="f4")
assert_sycl_queue_equal(y.sycl_queue, x.sycl_queue)
sycl_queue = dpctl.SyclQueue(device_y)
y = dpnp.astype(x, dtype="f4", device=sycl_queue)
assert_sycl_queue_equal(y.sycl_queue, sycl_queue)
Loading