Skip to content

Backport gh-1693 to maintenance/0.17.x branch #1694

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

Merged
merged 2 commits into from
May 23, 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
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ and complies with revision [2023.12](https://data-apis.org/array-api/2023.12/) o
* Fixed bug in basic slicing of empty arrays: [gh-1680](https://github.com/IntelPython/dpctl/pull/1680)
* Fixed bug in `tensor.bitwise_invert` for boolean input array: [gh-1681](https://github.com/IntelPython/dpctl/pull/1681)
* Fixed bug in `tensor.repeat` on zero-size input arrays: [gh-1682](https://github.com/IntelPython/dpctl/pull/1682)
* Fixed bug in `tensor.searchsorted` for 0d needle vector and strided hay: [gh-1694](https://github.com/IntelPython/dpctl/pull/1694)


## [0.16.1] - Apr. 10, 2024
Expand Down
27 changes: 17 additions & 10 deletions dpctl/tensor/libtensor/source/sorting/searchsorted.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -340,22 +340,29 @@ py_searchsorted(const dpctl::tensor::usm_ndarray &hay,
int simplified_nd = needles_nd;

using shT = std::vector<py::ssize_t>;

shT simplified_common_shape;
shT simplified_needles_strides;
shT simplified_positions_strides;
py::ssize_t needles_offset(0);
py::ssize_t positions_offset(0);

dpctl::tensor::py_internal::simplify_iteration_space(
// modified by refernce
simplified_nd,
// read-only inputs
needles_shape_ptr, needles_strides, positions_strides,
// output, modified by reference
simplified_common_shape, simplified_needles_strides,
simplified_positions_strides, needles_offset, positions_offset);

if (simplified_nd == 0) {
// needles and positions have same nd
simplified_nd = 1;
simplified_common_shape.push_back(1);
simplified_needles_strides.push_back(0);
simplified_positions_strides.push_back(0);
}
else {
dpctl::tensor::py_internal::simplify_iteration_space(
// modified by refernce
simplified_nd,
// read-only inputs
needles_shape_ptr, needles_strides, positions_strides,
// output, modified by reference
simplified_common_shape, simplified_needles_strides,
simplified_positions_strides, needles_offset, positions_offset);
}
std::vector<sycl::event> host_task_events;
host_task_events.reserve(2);

Expand Down
16 changes: 16 additions & 0 deletions dpctl/tests/test_usm_ndarray_searchsorted.py
Original file line number Diff line number Diff line change
Expand Up @@ -355,3 +355,19 @@ def test_out_of_bound_sorter_values():
p = dpt.searchsorted(x, x2, sorter=sorter)
# verify that they were applied with mode="wrap"
assert dpt.all(p == dpt.arange(3, dtype=p.dtype))


def test_searchsorted_strided_scalar_needle():
get_queue_or_skip()

a_max = 255

hay_stack = dpt.flip(
dpt.repeat(dpt.arange(a_max - 1, -1, -1, dtype=dpt.int32), 4)
)
needles_np = np.squeeze(
np.random.randint(0, a_max, dtype=dpt.int32, size=1), axis=0
)
needles = dpt.asarray(needles_np)

_check(hay_stack, needles, needles_np)
Loading