I am doing quite a bit of work with JAX (the Python library used here) in a high-performance numerical computing context. On GPU/TPU, it is not going to reach

I am doing quite a bit of work with JAX (the Python library used here) in a high-performance numerical computing context. On GPU/TPU, it is not going to reach perfect 100% hardware usage, but it is going to get close enough (far above vanilla Python performance) and be significantly more productive than alternatives. That makes it a sweet spot for research (where you will want to tweak things as you go) and extremely complex codes (where you already need to put your full focus on the correctness of the code). I highly recommend it to domain experts who need performance for their research project.