3 ms·
I think numba still makes sense for loopy algorithms but not so much if youre more vector oriented given that Jax is more or less a drop in replacement for nump
by usgroup 4y ago
I think numba still makes sense for loopy algorithms but not so much if youre more vector oriented given that Jax is more or less a drop in replacement for numpy and is shockingly fast.
- galangalalgol 4y agoI have used pytorch as a (almost) drop in replacement for numpy. Are there good reasons to look at jax instead assuming I'm doing DSP and not ML?
- martinsmit 4y agoIf you are doing array or vector-based work where the operations can be written as maps as opposed to for loops then JAX is king imo.
- patrickkidger 4y agoHonestly, the two are now incredibly close. JAX introduced a lot of cool concepts (e.g. autobatching (vmap), autoparallel (pmap)) and supported a lot of things that PyTorch didn't (e.g. forward mode autodiff). And at least for my applications (scientific computing), it was much faster (~100x) due to a much better JIT compiler and reduced Python overhead. ...but! PyTorch has worked hard to introduce all of the former, and the recent PyTorch 2 announcement was primarily about a better JIT compiler for PyTorch. (I don't think anyone has done serious non-ML benchmarks for this though, so it remains to be seen how this holds up.) There are still a few differences. E.g. JAX has a better differential equation solving ecosystem. PyTorch has a better protein language model ecosystem. JAX offers some better power-user features like custom vmap rules. PyTorch probably has a lower barrier to entry. (FWIW I don't know how either hold up specifically for DSP.) I'd honestly suggest just trying both; always nice to have a broader selection of tools available.