We were very heavy numba users at my former company. I would even go so far as to say numba was probably the biggest computational enabler for the product. I’ve also made a small contribution to the library. It’s a phenomenal library for developing novel computationally intensive algorithms on numpy arrays. It’s also more versatile than Jax. In presentations, I’ve heard Leland McInnes credits numba often when he spea…
> 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
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 built a numerical optimisation library with it, and although a few things became verbose, it was a rather pleasant experience as the natural vmapping made everything a breeze, I didn't have to write the gradients for my testing functions, except for special cases that involved exponents and logs that needed a bit of delicate care.