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.

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

Jump to bottom

Fix ExecutionPlacementError for dpnp.take_along_axis #1702

Merged
merged 4 commits into from
Feb 9, 2024
Merged

Fix ExecutionPlacementError for dpnp.take_along_axis #1702

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
240a759
Follow compute follows data to fill fancy_index
vlad-perevezentsev Feb 9, 2024
9430436
Update take_along_axis tests to cover the issue
vlad-perevezentsev Feb 9, 2024
11edc35
Merge master into fix_take_along_axis
vlad-perevezentsev Feb 9, 2024
e56be6f
Update test_take_along_axis
vlad-perevezentsev Feb 9, 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
7 changes: 6 additions & 1 deletion dpnp/dpnp_iface_indexing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -111,7 +111,12 @@ def _build_along_axis_index(a, indices, axis):
else:
ind_shape = shape_ones[:dim] + (-1,) + shape_ones[dim + 1 :]
fancy_index.append(
dpnp.arange(n, dtype=indices.dtype).reshape(ind_shape)
dpnp.arange(
n,
dtype=indices.dtype,
usm_type=indices.usm_type,
sycl_queue=indices.sycl_queue,
).reshape(ind_shape)
)

return tuple(fancy_index)
Expand Down
36 changes: 32 additions & 4 deletions tests/test_sycl_queue.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1520,28 +1520,56 @@ def test_clip(device):
assert_sycl_queue_equal(x.sycl_queue, y.sycl_queue)


@pytest.mark.parametrize("func", ["take", "take_along_axis"])
@pytest.mark.parametrize(
"device",
valid_devices,
ids=[device.filter_string for device in valid_devices],
)
def test_take(func, device):
def test_take(device):
numpy_data = numpy.arange(5)
dpnp_data = dpnp.array(numpy_data, device=device)

dpnp_ind = dpnp.array([0, 2, 4], device=device)
np_ind = dpnp_ind.asnumpy()

result = getattr(dpnp, func)(dpnp_data, dpnp_ind, axis=None)
expected = getattr(numpy, func)(numpy_data, np_ind, axis=None)
result = dpnp.take(dpnp_data, dpnp_ind, axis=None)
expected = numpy.take(numpy_data, np_ind, axis=None)
assert_allclose(expected, result)

expected_queue = dpnp_data.get_array().sycl_queue
result_queue = result.get_array().sycl_queue
assert_sycl_queue_equal(result_queue, expected_queue)


@pytest.mark.parametrize(
"data, ind, axis",
[
(numpy.arange(6), numpy.array([0, 2, 4]), None),
(
numpy.arange(6).reshape((2, 3)),
numpy.array([0, 1]).reshape((2, 1)),
1,
),
],
)
@pytest.mark.parametrize(
"device",
valid_devices,
ids=[device.filter_string for device in valid_devices],
)
def test_take_along_axis(data, ind, axis, device):
dp_data = dpnp.array(data, device=device)
dp_ind = dpnp.array(ind, device=device)

result = dpnp.take_along_axis(dp_data, dp_ind, axis=axis)
expected = numpy.take_along_axis(data, ind, axis=axis)
assert_allclose(expected, result)

expected_queue = dp_data.get_array().sycl_queue
result_queue = result.get_array().sycl_queue
assert_sycl_queue_equal(result_queue, expected_queue)


@pytest.mark.parametrize(
"device",
valid_devices,
Expand Down
26 changes: 26 additions & 0 deletions tests/test_usm_type.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -570,6 +570,32 @@ def test_take(func, usm_type_x, usm_type_ind):
assert z.usm_type == du.get_coerced_usm_type([usm_type_x, usm_type_ind])


@pytest.mark.parametrize(
"data, ind, axis",
[
(numpy.arange(6), numpy.array([0, 2, 4]), None),
(
numpy.arange(6).reshape((2, 3)),
numpy.array([0, 1]).reshape((2, 1)),
1,
),
],
)
@pytest.mark.parametrize("usm_type_x", list_of_usm_types, ids=list_of_usm_types)
@pytest.mark.parametrize(
"usm_type_ind", list_of_usm_types, ids=list_of_usm_types
)
def test_take_along_axis(data, ind, axis, usm_type_x, usm_type_ind):
x = dp.array(data, usm_type=usm_type_x)
ind = dp.array(ind, usm_type=usm_type_ind)

z = dp.take_along_axis(x, ind, axis=axis)

assert x.usm_type == usm_type_x
assert ind.usm_type == usm_type_ind
assert z.usm_type == du.get_coerced_usm_type([usm_type_x, usm_type_ind])


@pytest.mark.parametrize(
"data, is_empty",
[
Expand Down
Loading