[time_to_smile]

What I love about JAX is that it essentially just makes Python into a performant, differentiable programming language.

I'm a pretty big fan of moving away from thinking about ML/Stats/etc specifically and people should more generally embrace the idea of differentiable programming as just a way to program and solve a range of problems.

JAX means that the average python programmer just needs to understand the basics of derivatives and their use (not how to compute them, just what they are and why they're useful) and suddenly has an amazing amount of power they can add to normal code.

The real power of JAX, for me at least, is that you can write the solution to your problem, whatever that problem may be, and use derivatives and gradient descent to find an answer. Sometimes this solution might be essentially a neural network, other times the generalized linear model, but sometimes it might not fit obviously into either of these paradigms.

[adgjlsfhk1]

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.