Jake Vanderplas

We're blocked by two `jax.experimental.jet` issues: - https://github.com/google/jax/issues/12263 - missing jet rule for `dynamic_update_slice_p`

Added a second commit with benchmarks to prevent regression. Before: ``` name cpu/op bench_repeated_static_indexing 11.1s ±37% bench_repeated_static_slicing 9.68s ±39% name time/op bench_repeated_static_indexing 11.2s ±38% bench_repeated_static_slicing 9.72s ±40% name allocs/op bench_repeated_static_indexing...

Hi - I suspect this may be related to the general errors people have been seeing on Colab TPU for the last few weeks; see https://github.com/googlecolab/colabtools/issues/3009 It may be that...

This should be fixed in the most recent jax & jaxlib release. Thanks!

embed options do not affect the chart specification, they control the HTML and Javascript code used to display the chart specification in the frontend. That's why we need a weird...

You can check the effect like this, for example: ```python >>> import altair as alt >>> import pandas as pd >>> df = pd.DataFrame({"x": [1, 2, 3, 4], "y": [1,...

Sorry this has caused problems for you! The out-of-bounds indexing issue is discussed in the documentation here: https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#out-of-bounds-indexing Do you have suggestions for a better place to surface the issue...

I like the idea of having a new dedicated doc for performance tips and pitfalls

Regarding reworking the Sharp Bits doc, I recently added a section on miscellaneous divergences between numpy and JAX. It might be nice to rearrange things so all the differences between...

@nalzok - there is currently an implementation of this, but only for TPU. See https://github.com/google/jax/tree/main/jax/experimental/compilation_cache for details, and https://github.com/google/jax/issues/2490 where this kind of request is tracked.