I have definitely gotten significant speed ups using Jax. I try to vectorize my code as much as possible but that’s not always feasible for every aspect of an a

I have definitely gotten significant speed ups using Jax. I try to vectorize my code as much as possible but that’s not always feasible for every aspect of an algorithm. Even for the vectorized operations, the JIT compilation gets me a speed up when used correctly (I.e. structuring code to minimize recompiling). Jax also has the added bonus of trivially running your algorithm on the GPU.