Skip to content

Commit 28f86cc

Browse files
committed
ENH: searchsorted: allow python scalars for x2
cross-ref data-apis/array-api#982
1 parent ce93ed2 commit 28f86cc

File tree

1 file changed

+45
-0
lines changed

1 file changed

+45
-0
lines changed

array_api_tests/test_searching_functions.py

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -291,3 +291,48 @@ def test_searchsorted(data):
291291
except Exception as exc:
292292
exc.add_note(repro_snippet)
293293
raise
294+
295+
296+
### @pytest.mark.min_version("2025.12")
297+
@given(data=st.data())
298+
def test_searchsorted_with_scalars(data):
299+
# 1. draw x1, sorter and side exactly the same as in test_searchsorted
300+
x1_dtype = data.draw(st.sampled_from(dh.real_dtypes))
301+
_x1 = data.draw(
302+
st.lists(
303+
xps.from_dtype(x1_dtype, allow_nan=False, allow_infinity=False),
304+
min_size=1,
305+
unique=True
306+
),
307+
label="_x1",
308+
)
309+
x1 = xp.asarray(_x1, dtype=x1_dtype)
310+
if data.draw(st.booleans(), label="use sorter?"):
311+
sorter = xp.argsort(x1)
312+
else:
313+
sorter = None
314+
x1 = xp.sort(x1)
315+
316+
kw = data.draw(hh.kwargs(side=st.sampled_from(["left", "right"])))
317+
318+
# 2. draw x2, a real-valued scalar (IOW, an int or a float)
319+
x2 = data.draw(hh.scalars(st.sampled_from([xp.int32, xp.float64]), finite=True))
320+
321+
# 3. testing: similar to test_searchsorted, modulo `out.shape == ()`
322+
repro_snippet = ph.format_snippet(
323+
f"xp.searchsorted({x1!r}, {x2!r}, sorter={sorter!r}, **kw) with {kw = }"
324+
)
325+
try:
326+
out = xp.searchsorted(x1, x2, sorter=sorter, **kw)
327+
328+
ph.assert_dtype(
329+
"searchsorted",
330+
in_dtype=[x1.dtype], #, x2.dtype
331+
out_dtype=out.dtype,
332+
expected=xp.__array_namespace_info__().default_dtypes()["indexing"],
333+
)
334+
# TODO: values testing
335+
ph.assert_shape("searchsorted", out_shape=out.shape, expected=())
336+
except Exception as exc:
337+
exc.add_note(repro_snippet)
338+
raise

0 commit comments

Comments
 (0)