I just love pythons ease of use, and the few times I've tried Julia it has felt clanky in comparison, and then Jax just makes everything so amazingly fast.
"Clanky"? That is a word I would use when comparing Julia and Python, but I would reverse the roles. I mean, python works well, and has almost everything, but it really feels, well, clanky.
I'm currently working on a re-write of a small model from JAX to Julia and finding the Julia code so much easier to write, it's more concise, and find the debugging tools easier to work with in Julia.
for problems where Jax works well, in theory there could be a great Julia equivalent, but in practice one has not emerged. instead every 3 or so years, a brilliant solo researcher writes an even better autograd library which is ultimately abandoned.
the language is at its best in other areas. for a few things (differential equations, mathematical optimization) it is actually state of the art.
Even as a full-blown Julia hater, it exists at a really nice level of abstraction for working in differential equations and Chris Rackaukas has written some fantastic software in that area. It can be a great tool for some domains (ODEs, statistical modeling/visualization) and I can see why Pumas uses the language.
I don’t think that Julia is a great tool to replace the complicated physics simulation code that would have traditionally been written in C++/Fortran, I think for that you’ll generally be better off using JAX/Rust/C++.
Sometimes it really feels like JAX should be a more general DSL that can embed in other languages, as long as you have some way to turn structures into pytrees and functions into jaxpr’s.
I just love pythons ease of use, and the few times I've tried Julia it has felt clanky in comparison, and then Jax just makes everything so amazingly fast.