Thanks for your comment - in the JIT section under "Jax Transformations" I implement the same calculation on NP and JAX on CPU (and TPU) to get a direct compari

Thanks for your comment - in the JIT section under "Jax Transformations" I implement the same calculation on NP and JAX on CPU (and TPU) to get a direct comparison for CPU. I mentioned in another comment that the calculation lends itself well to being JITted, but I have some preliminary data for a follow-up JAX vs NumPy performance comparison. It's clear so far that NumPy can outperform JAX on CPU for small computations (unsurprisingly). Please let me know if you'd be interested in seeing a more thorough analysis!