Make `xp.searchsorted` accept Python scalars for the second argument

[ogrisel]

import array_api_strict as xp

a = xp.arange(10, device=xp.Device("device1"))
xp.searchsorted(a, 42)

raises:

Traceback (most recent call last):
  Cell In[5], line 5
    xp.searchsorted(a, 42)
  File ~/miniforge3/envs/dev/lib/python3.13/site-packages/array_api_strict/_flags.py:395 in wrapper
    return func(*args, **kwargs)
  File ~/miniforge3/envs/dev/lib/python3.13/site-packages/array_api_strict/_searching_functions.py:78 in searchsorted
    if x1.dtype not in _real_numeric_dtypes or x2.dtype not in _real_numeric_dtypes:
AttributeError: 'int' object has no attribute 'dtype'

This is a bit annoying as it requires to instead write: xp.searchsorted(a, xp.asarray(42, device=a.device)) which feels unnecessarily verbose.

Even PyTorch accepts the following without complaining, for instance:

import array_api_compat.torch as xp

a = xp.arange(10, device="mps")
xp.searchsorted(a, 42)

However, the SPEC does not mention Python scalar support explicitly, so maybe it would need to be updated first?

https://data-apis.org/array-api/latest/API_specification/generated/array_api.searchsorted.html#searchsorted

[lucascolley]

I guess this is similar to #860

[lucascolley]

This was at least a conscious decision originally, #699 (comment)

[ogrisel]

requires that x2 be an array. CuPy and TensorFlow require that x2 be an array, while NumPy, PyTorch, and others allow x2 to be a scalar. In this PR, I've chosen to restrict x2 to be an array. For the scalar use case, downstream consumers can wrap scalar values as zero-dimensional arrays.

Indeed, thanks for looking this up. Still, I think that calling searchsorted with a scalar as second argument is quite common, and it's easy to forget about passing the device argument when developing with numpy on a CPU and then get the code crash when running with another on GPUs.

I think we could make the spec evolve and handle the discrepancy in scalar support via array_api_compat.

[lucascolley]

Yes, I think I am +1 because the usage feels similar to xp.where.

[kgryte]

I'll go ahead and add this to the next array API workgroup meeting as a discussion topic.