4 ms·

> 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

> 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

fasttriggerfish 4y ago
I’ve been disappointed with Jax which I was trying to use for backward auto differentiation. The issue is that XLA JIT compilation is very slow and easily adds half a minute of overhead to the first call of the base function just by using jax.numpy instead of numpy, which made it a non starter for my use case. It’s definitely optimised for large flow computations where the JIT overhead is dwarfed by the rest. In the end I reverted to using autograd which did the job fine. I had never heard of tai chi until now, I’m curious how it compares.
PartiallyTyped 4y ago
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 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. [1] https://github.com/HIPS/autograd https://github.com/HIPS/autograd
- rg111 4y ago
  > You may be able to get the equivalent of jax via numba+numpy+autograd[1], but I haven't tried it before. I have tried it before [0]. And it works. As expected. Not much friction. [0]: sorry no github

fasttriggerfish 4y ago

I’ve been disappointed with Jax which I was trying to use for backward auto differentiation. The issue is that XLA JIT compilation is very slow and easily adds half a minute of overhead to the first call of the base function just by using jax.numpy instead of numpy, which made it a non starter for my use case. It’s definitely optimised for large flow computations where the JIT overhead is dwarfed by the rest. In the end I reverted to using autograd which did the job fine. I had never heard of tai chi until now, I’m curious how it compares.

PartiallyTyped 4y ago

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 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. [1] https://github.com/HIPS/autograd https://github.com/HIPS/autograd

rg111 4y ago
> You may be able to get the equivalent of jax via numba+numpy+autograd[1], but I haven't tried it before. I have tried it before [0]. And it works. As expected. Not much friction. [0]: sorry no github