I mean, not really? Most Jax code is only marginally more restrictive than numpy, and most of the annoyances I have with Jax are shared with numpy.

And having spent too much of my life writing some extremely high performance physics code, 99% of the code complexity comes from the stupid stuff, like handling configuration and IO and restarting. Python handles that stuff great. Oh, and tests. Ever tried to write unit tests in Fortran? You'll hug python tight and promise it you'll never leave again.