ENH: add axis support to np.searchsorted - #32346
Conversation
Sorry, something went wrong.
|
We talked about this during the triage review meeting. We think that maybe the simplest thing is to simply drop the axis keyword argument, keep the gufunc, and force people who would need to use axis to reshuffle. |
Sorry, something went wrong.
|
In the last commit I've droppied axis and I'm only searching along axis=-1 by default internally. There are unrelated CI failures due to cython. I have not updated the PR description and title. I just pushed this change just to brainstorm what's best. |
Sorry, something went wrong.
Signed-off-by: Iason Krommydas <iason.krom@gmail.com>
Signed-off-by: Iason Krommydas <iason.krom@gmail.com>
Signed-off-by: Iason Krommydas <iason.krom@gmail.com>
Signed-off-by: Iason Krommydas <iason.krom@gmail.com>
Signed-off-by: Iason Krommydas <iason.krom@gmail.com>
Signed-off-by: Iason Krommydas <iason.krom@gmail.com>
Signed-off-by: Iason Krommydas <iason.krom@gmail.com>
Signed-off-by: Iason Krommydas <iason.krom@gmail.com>
|
Thanks, this is looking nice, the implementation is coming along really well and many general improvements! Just a few high level comments for now, just opening them up for discussion really in any case anyone has any thoughts. From a brief look my biggest concern is about having four gufuncs for the same operation. That seems like we are exposing implementation/selection details to the gufunc object itself. I think it is better to keep it to one gufunc and use the void *data argument (which seems to be currently perm) to select among implementations for sides and sorted. I think this is at least true for sides, which is very much a parameter. This could be implemented as a struct in the same way as for example PyArrayMethod_SortParameters. This getter-style architecture also lets us evolve arguments if we ever need to. FWIW, I think a similar discussion was actually had in #29737, but for algorithm choice. The other point is that we are still wrapping the gufunc, with PyArray_SearchSorted_int? In that case could it not be better to redirect from PyArray_SearchSorted itself? It seems here we have something like PyArray_SearchSorted_int -> gufunc (if found) -> otherwise PyArray_SearchSorted, but PyArray_SearchSorted -> PyArray_SearchSorted_int (or directly use the gufunc) -> otherwise fallback seems more intuitive to me. Of course it's consistent with the sorts that way as well, but I guess I'm just wondering if the entrypoint here should be changed, to me it seems this adds a layer in between of the existing two layers, it's not really a change in the entrypoint in essence. |
Sorry, something went wrong.
I don't immediately see how sorter can be in the same gufunc loop. It's an entirely new array and changes the ufunc signature. It looks like we need at least two to me. Regarding the side it can definitely be a scalar (() in the signature) input to the loop however I thought more of it as static and not something I wanted to expose to the loop even with the void *data. The latter is possible though.
PyArray_SearchSorted is public C api. It is not allowed to change so I had to keep it intact. This is why I resorted to a new PyArray_SearchSorted_int which is not exposed publicly. The point is that PyArray_SearchSorted should remain as is (only did internal refactoring inside the function but did not break api here). You are right that what you say seems more intuitive and it is usually how it works in numpy with functions that have an internal (_int) function too but here we can't do that cause It'd break the api. |
Sorry, something went wrong.
PR summary
Closes #4224.
np.searchsorted and ndarray.searchsorted now take an axis keyword and no longer require the sorted array to be one-dimensional. Each 1-D slice along axis (the last by default, or the flattened array when axis=None) is searched independently.
Internally this is four private gufuncs, _searchsorted_{left,right} with signature (n),(m?)->(m?) and _searchsorted_{left,right}_sorter with (n),(m?),(n)->(m?). side picks the gufunc statically, so the kernels in npysort/binsearch.cpp are unchanged and still selected as get_binsearch_func(dtype, side) did.
(n),(m?)->(m?) is the signature and what array-api-extra implements, with the flexible m? so a scalar v keeps returning a scalar. Keys are v's last axis and leading axes broadcast against a's other dimensions. The searched dimension of the result sits at axis, so a of shape (2,3,4) with axis=0 and 5 keys gives (5,3,4). Regarding Jake's question, the gufunc default would put query dims last here, but axes is used to place them at axis instead, matching fft, take_along_axis and sort. So this differs from the JAX layout. There is no batch_dims equivalent, leading dims broadcast right-aligned, so the batched cases need a newaxis manually from the user.
The gufunc is an internal layer rather than the Python entry point: dispatch stays in C, PyArray_SearchSorted is unchanged and serves as the fallback, and the legacy ArrFunc_Compare layer is untouched. The dtype slot is not here, but a registered loop is the only thing that grants a dtype N-D support, which is the seam a searchsorted_meth slot plugs into, so that follow-up is additive and would close the existing TODO in item_selection.c.
axis works for all builtin dtypes and for run-time registered ones via a promoter. StringDType stays 1-D for now. Operands defining __array_ufunc__ are consulted.
AI Disclosure
AI has been used a lot in brainstorming about this PR about the signature and potential ways to implement the PR and to create quick rough sketches to try out myself which would have taken me many days to do manually. Eventually I iterated with AI to figure out how to logically do the dtype handling and for review.