Earlier quoted context omitted.
> It’s also more versatile than Jax Does numba do automatic differentiation? I view JAX as primarily an automatic differentiation tool with the bonus that it makes great use of XLA and can easy make use of GPU/TPUs. I don’t usually see numba and JAX as solving the same problem, but would be excited to be wrong
They solve different problems. Numba compiles functions down to machine code or cuda kernels, that's it. XLA is "higher level" than what Numba produces. You may be able to get the equivalent of jax via numba+numpy+autograd[1], but I haven't tried it before. IMHO, jax is best thought of as a numerical computation library that happens to include autograd, vmapping, pmapping and provides a high level interface for XLA.…
I have tried it before [0]. And it works. As expected. Not much friction.
[0]: sorry no github