even_spacing option to omit binary search#148
Draft
YigitElma wants to merge 3 commits intof0uriest:mainfrom
Draft
even_spacing option to omit binary search#148YigitElma wants to merge 3 commits intof0uriest:mainfrom
even_spacing option to omit binary search#148YigitElma wants to merge 3 commits intof0uriest:mainfrom
Conversation
YigitElma
commented
Dec 23, 2025
| i = jnp.clip(jnp.searchsorted(x, xq, side="right"), 1, len(x) - 1) | ||
| else: | ||
| dx = x[1] - x[0] | ||
| i = jnp.clip(jnp.floor((xq - x[0]) / dx).astype(int) + 1, 1, len(x) - 1) |
Author
There was a problem hiding this comment.
Check if this still works for periodic inputs
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
For evenly spaced regular grids, the neighboring points of the query positions can be found without a binary search. The benefits of this are hopefully twofold:
diffrax, due tojax.numpy.searchsortedeach kernel needs all the arrays and hence different kernels have to wait for another to finish. Ideally, with this change, each kernel only needs to accessThe current version is working fine and I think would benefit others too, however, for the second point, I need to make more profiling, and see if it is making a meaningful difference or needs additional changes.