Repository navigation
searchsorted should get a axis kwarg #4224
Description
Is there a technical reason why searchsorted only works with 1d arrays? In a current problem, I have an array v.shape = (61, 160, 320) and an array a.shape = (1, 160, 320). It would be great if I could do np.searchsorted(a, v, axis=0) to yield a (1, 160, 320) array of the searched indices.
Activity
charris commentedon Feb 15, 2014
@jaimefrio Any thoughts about this?
jaimefrio commentedon Feb 15, 2014
I think it makes a lot of sense. I have a PR in the works that adds an
axisargument tonp.bincount, which is being my crash course onNpyIterand nested iteration. Much of what I am doing there should be relatively easy to apply tonp.searchsorted, especially now that the new binary search functions can take strided arrays for all three arguments.There may be a few quirks to iron out with the UI: broadcasting mixed with an axis argument can get tricky and confusing very fast. I'll ping the list when I have looked into it with a little more care.
andreas-h commentedon Mar 24, 2015
@jaimefrio any news relevant to this issue?
jaimefrio commentedon Mar 24, 2015
Unfortunately there aren't many good news. We have discussed this elsewhere, but my PR for
bincountderailed (#4330) because, even though it worked, it added over 600 lines of code to a function that is barely over 300 right now. So the way to go doesn't seem to be to just multiply the size of every function that could use anaxisargument by 3, but rather to build a middle layer that abstracts the iterator complexity, and can be reused in all those functions.I might try to get this started in the next few days by sending a longish proposal to the mailing list.
AlJohri commentedon Jan 19, 2016
just dropping a note that I came across a need for the axis argument as well. want to perform searchsorted by row over a large matrix
thanks for your hard work! 👍
rob-smallshire commentedon May 18, 2017
+1 on a need for this. Major performance bottleneck in my application.
SreenivasVRao commentedon Apr 14, 2018
Any progress with this? I'm also looking for an axis argument to
searchsortedaltaetran commentedon Apr 21, 2018
Definitely looking for the same
eric-wieser commentedon Apr 21, 2018
@mhvk: This might be a good place to expose the new gufunc axis arguments
mhvk commentedon Apr 22, 2018
It would be lovely to expose gufunc arguments, but I don't think the
binsearchfunctions are gufuncs, or are they?charris commentedon Apr 23, 2018
Everything in
numpy/core/src/npysortis put into anpysort.alibrary, which is then linked with multiarray. It should not be too difficult to link them into gufuncs, but the sorting functions at least require contiguous inputs. I've thought of adding astepargument, but suspect that there is a tradeoff between speed of access and the overhead of making a copy. Having a gufunc is not quite the same as having an arbitrary axis, but would allow working with stacks of vectors. One difficulty would be if we wanted to broadcast the keys. The output dimensions would also be determined by the keys, which would usually not have dimensions compatible with the searched data, so that could be interesting.eric-wieser commentedon Apr 23, 2018
It is now that @mhvk allowed gufuncs to specify axis arguments
Isn't this just
(i),(j)->(i)?20 remaining items
MaanasArora commentedon Jul 23, 2026
Still have to catch up on the whole discussion, but I think Marten mostly meant ArrayMethods in the dtype sense - I don't think that precludes the use of gufuncs per se! In fact the (arg)partition spec does take the
kthindex as an additional array argument exactly as in a gufunc. We do some processing on the ktharray manually, so it is true that the broadcasting logic may be more gufunc-native forsearchsorted.For consistency, we could use very similar dtype dispatching logic as in sorting / partitions, instead of relying entirely on a gufunc as the entrypoint, which does not mean we cannot use a gufunc entirely under the hood (in fact I think we should, to the extent possible). That would include having slots on the dtype and keeping the
PyArray_SearchSortedfunction logic we currently have. It will let us not touch the legacy layer in a way that mirrors the sorts (which also use theArrFunc_Compareslot). I think could be nicely flexible, both on the legacy and interoperability end, to keep gufuncs as an internal implementation layer rather than a direct Python layer (or even do both and interleave in C if you like having a gufunc anyway).mhvk commentedon Jul 24, 2026
@ikrommyd - just a quick note on the shapes - with
np.atleast_1d(), I think you'd indeed need to squeeze the last axis if a scalar was passed in, if only to remain completely backwards compatible.@jakevdp - with the
gufuncaxesargument, one has full control over which axis of each input get used, and where the axis gets put in the output. As noted, for numpy by default these are always the last axis, so it would be[-1, -1, -1]in this case. EDIT: so I think one can cover all cases you listed EDIT except the multiple dimensions in one go.What we do need to decide, is what
axiswould mean. In thegufuncsignature, it is only allowed if there is only a single shared dimension for all arguments (say,(n),(n)->()), while in this case there would at least have to be an exception for scalarv.ikrommyd commentedon Jul 24, 2026
I don't see how one covers this case from jax with
axesautomatically. It needs 3 loop dims while operands give at most 2 no? I may be stupid here thoughYou can achieve it with some python wrapping creating the new dimensions.
mdhaber commentedon Jul 24, 2026
I haven't thought about this for a while, but I appreciate that it's happening. I'd prefer that the behavior of array-api-extra's
search sortedbe the default and suggest that users can reshape as needed to make it work without additional arguments, but as long as the array-api-extra behavior is achievable, I'm happy.mhvk commentedon Jul 24, 2026
@ikrommyd - you're right, one cannot do that case without a reshape. I think we've at some point discussed expanding a single axis to multiple ones (as can be done with
np.sum(a, axis=(2,3))but that is not possible forgufuncs.mdhaber commentedon Aug 15, 2026
@ikrommyd Thought I'd check in. How is this going?
I wanted to mention more specifically that the
(n),(m?),->(m?)behavior is exactly what we did in array-api-extra's searchsorted in hope of getting other libraries and the standard to adopt it.1 I've worked a lot with gufuncs (e.g. added batch support to all ofscipy.linalg) and I think it is the natural choice.Adding an
axesargument to NumPy might be fine, but I wouldn't let interface concerns about it hold things up. Treating the last axis of each array as the core dimension is sufficient, and I doubtaxeswould end up in the first pass of the standard, since there's not a ton of precedent for that keyword.Footnotes
The only difference is that we required the dimensions other than the last to be the same shape before calling the function. However, the broadcasting behavior is well defined and great to include in NumPy. I can't remember exactly why we didn't include it, but it's probably to avoid having to document how it works (e.g. remove core shapes, broadcast batch shapes, prepend to core shapes) and IIRC because there are other functions like it in which the Array API could broadcast but doesn't (e.g.
concat). ↩ikrommyd commentedon Aug 18, 2026
I waited a bit in case people had more opinions and I had also sent a message to the mailing list but I'm working on an implementation this week.
lucascolley commentedon Aug 20, 2026
(docs at https://data-apis.org/array-api-extra/generated/array_api_extra.searchsorted.html)
ikrommyd commentedon Aug 21, 2026
So in the last numpy triage meeting, we discussed considering not exposing an axis argument but going the same route as the array_api_extra and doing it along the last axis and having the user be responsible of swapping the axes if they want. What do you people think of that approach? I'm mainly asking @mdhaber here because you mentioned earlier that this would unblock adding an axis argument in scipy stats functions. Also the original issue from 2014 here is actually asking for searching along axis=0.
I'm just trying to figure out what use cases are there and whether last axis only enables all of them and whether there is actually a need for exposing axis.
mhvk commentedon Aug 21, 2026
I think as long as we use a
gufuncunderneath, I'm fine. If we cannot yet match even the fanciest batching, that seems OK.ikrommyd commentedon Aug 21, 2026
Yeah a gufunc would be underneath and that's good for user dtypes too as they can register the loop to support ND searchsorted. I'm only talking about exposing the axis kwarg or forcing it internally to be axis=-1 basically. If axis=-1 is like 99% of the use cases, we may not want to expose it for example. The core work of turning it into a gufunc is in my draft PR already.
MaanasArora commentedon Aug 21, 2026
I agree using a gufunc underneath is sufficient for now; I'd just design the axis semantics such that we can eventually expose
axis=to Python as well. Based on what Matt said above about SciPy use-cases it already seems clearly worth it to me. I don't mind doing a little reshaping on the Python end for legacy compatibility eventually if needed. Also, this helps divide up the PRs for easier review!Edit: I mean that, by looking at many
scipy.interpolate,scipy.integratefunctions, and other modules,axis=is passed as an argument. (See interp1d for example but there are many more of course!) I would imagine it would be annoying to swap axes to work withsearchsortedfor each, but Matt could speak better to that. Even if it's not particularly annoying here this is quite a bit of precedent thataxis=can be useful.mdhaber commentedon Aug 21, 2026
Going without
axis/axesis absolutely fine. Thanks!ikrommyd commentedon Aug 25, 2026
@mdhaber does that mean you'd do
moveaxisshenanigans in scipy or does axis=-1 cover all the use-cases you want? 🤣 Cause I care about the difference a lot.mdhaber commentedon Aug 25, 2026
axis=-1covers all the use cases that can be had. A standardaxisargument is always redundant with other features (namelymoveaxisand, for tuple axes ,reshape) and provided for convenience.