Jake Vanderplas

One other thing I just noticed: everything in the body of `for r in range(n_trials):` looks to be JIT-compatible. You might think about wrapping that in a function and compiling...

Ah, I didn't realize you were running on GPU. In that case, the notes about sequential updating blocking asynchronous execution are even more important! I would start by trying to...

Dispatch means executing a JAX function (i.e. sending a set of data and instructions to XLA) For example, this has 100 dispatch events: ```python x = jnp.arange(100) def f(x): out...

@dcalacci there's not much to debug unfortunately: the metal plugin is still experimental and very incomplete, and so you should expect to run into these kinds of issues when using...

Failure looks like it may be a flake; I'll leave this open for a day to see if it repeats.

Sorry, I'm having trouble working out what your functions are actually doing (long day...) could you show an example of a simple call that fails that you think should not...

Thanks - edited to make the result specification more canonical. @sharadmv, can you take a look?

We did a few PRs related to this recently: #17139, #17307, #17350, #17651, #17666, #17983, #18102, #18107, #18156 There's still some work to do, but I think we're beyond the...

Reopening, since there are a few more big ones to do (`jax.core`, `jax.interpreters`)

What happens if you wrap the jax function in `jax.jit`? Also, can you include details on how you ran the benchmarks? Keep in mind these tips to make sure you're...