This isn't quite true. Jax works well for "quasistatic" code, but can't handle more dynamic types of problems (see https://www.stochasticlifestyle.com/useful-algorithms-that-a... for a more detailed explanation).

Jax is definitely the right direction for the python ecosystem, but it can't solve all your problems. At some point you still need a fast language.