I love PyTorch, but my experience with jits embedded in Python (eg. Numba) has been everything but simple, nevermind trivial. I'll really have to try it to believe it.
I’ve had the opposite experience with numba in production. It works almost flawlessly, very easy to reason about the generated code and inspect annotations, easy to debug.
Do you interact with numpy or other compiled numerical packages? this is where it usually breaks for me. The thing is that I use numpy, scipy, keras, tensorflow, etc. in literally every project, making numba not too useful
That’s very cool. Numba and Cython work extremely well with virtually no overhead or extra effort on my part, so jax doesn’t seem like it would buy me much for most of my work. But I can imagine a lot of projects where jax woukd be useful, and I plan to keep current on best practices for it.