Jax is great, however you are still in the python ecosystem right? So, as you requirements grow, inevitably you will be at a point where 10% of your app goes through 'regular' python and if you don't "drop to C" - making this a two language problem - you will find your app in now 10x slower than an equivalent fast language. It is unbelievable how slow regular python (not using NumPy etc..) is.