Skip to content

Fix tensor.searchsorted for x1 with strides and scalar (0D) x2 #1693

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
May 22, 2024
Merged

Fix tensor.searchsorted for x1 with strides and scalar (0D) x2 #1693

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
97b1c82
Adds code to handle edge case of strided input and scalar `needle` in…
ndgrigorian May 22, 2024
8c6a57e
Adds a test for fix to gh-1689
ndgrigorian May 22, 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
27 changes: 17 additions & 10 deletions dpctl/tensor/libtensor/source/sorting/searchsorted.cpp
View file
Open in desktop
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
View file
Open in desktop
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