NumPy can't run on accelerators, which is part of where the difference stems from. I implemented the same operation first with both NP/JAX on CPU, then with JAX

NumPy can't run on accelerators, which is part of where the difference stems from. I implemented the same operation first with both NP/JAX on CPU, then with JAX on TPU and NP again on CPU, not sure if you saw that part. Obviously, the ability to use an accelerator makes JAX faster, but even without that on CPU it was faster. This is in part because of JIT, and in fairness the calculation in question does lend itself well to being expedited by JIT. I actually have some preliminary benchmarks for a follow up specifically on just NumPy vs JAX, and it has become clear so far that NumPy is better in certain cases, especially for small operations where the overhead of JAX is not worth it. In the article I mention this briefly, along with how JAX hasn't been focused on being optimized on CPU because they have bigger fish to fry, so to speak. I also link to the JAX documentation that has some comments comparing the two! Relating to TF - I don't actually use TF at any point, but I did use PyTorch for Hessian calculation. TF and PT obviously do both work on GPU, but JAX has the benefit of being able to JIT more and implement everything in terms of XLA (although TF obviously has XLA support as well, and PT kind of does but just to get PT working on TPU). Thinking about doing another article on a direct comparison of JAX with PT and TF - let me know if that's something you'd like to see!