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.
filter_custom_{jvp,vjp}
. This also brings with it an API tweak for backward compatibility. (E.g.filter_custom_jvp.defjvp -> filter_custom_jvp.def_jvp
for the new behaviour.)lax.select
ineqxi.while_loop
when not vmap'ing. This should improve speed slightly when not vmap'ing.eqxi.while_loop
buffers, in particular avoiding Big performance discrepancy between JAX and TensorFlow with in-place updates jax-ml/jax#10197. In particular this can be used to reduce some programs with quadratic runtime down to just linear runtime.