3 ms·

Can someone comment more on what makes JAX that much better for differentiable simulations than PyTorch? I'm working on a new module for work and none of my co

by uoaei 4y ago

Can someone comment more on what makes JAX that much better for differentiable simulations than PyTorch? I'm working on a new module for work and none of my colleagues have much experience developing ML per se. I'm trying to decide whether to force their hand by implementing v1 in PyTorch or JAX and differentiable physics simulations is a likely future use case. Why is PyTorch harder?

whimsicalism 4y ago
Because the `jax.numpy` operations & primitives are almost 1:1 with numpy, many working scientists who already have experience working with numpy will be able to figure out jax faster. It is also easier to rewrite existing code/snippets (say you were working on a non-differentiable simulator before) into jax if you already have them in numpy then to do the whole rewrite in pytorch. I will say that I think pytorch has improved its numpy compatability a lot in recent years, functions that I was convinced didn't exist with pytorch (like eigh) apparently actually do.
chazeon 4y ago
I have seen JAX-MD[1] but not sure about "much better". On the other hand, there is just no MD implemented with PyTorch. [1]: https://github.com/jax-md/jax-md https://github.com/jax-md/jax-md
- new_user55 4y ago
  There is [torch-md](https://github.com/torchmd/torchmd https://github.com/torchmd/torchmd)
patrickkidger 4y ago
At least prior to this announcement: JAX was much faster than PyTorch for differentiable physics. (Better JIT compiler; reduced Python-level overhead.) E.g for numerical ODE simulation, I've found that Diffrax (https://github.com/patrick-kidger/diffrax https://github.com/patrick-kidger/diffrax) is ~100 times faster than torchdiffeq on the forward pass. The backward pass is much closer, and for this Diffrax is about 1.5 times faster. It remains to be seen how PyTorch 2.0 will compare, of course! Right now my job is actually building out the scientific computing ecosystem in JAX, so feel free to ping me with any other questions.
- adgjlsfhk1 4y ago
  If you care about performance of differential physics you shouldn't use python. Diffrax is almost OKish, but is missing a ton of features (e.g. good stiff solvers, arbitrary precision support, events for anything other than stopping the simulation, ability to control the linear solve which are needed for large problems). For simple cases it can come close to the C++/Julia solvers, but for anything complicated, you either won't be able to formulate the model, or you won't be able to solve it efficiently.
  - patrickkidger 4y ago
    > If you care about performance This definitely isn't true. On any benchmark I've tried, JAX and Julia basically match each other. Usually I find JAX to be a bit faster, but that might just be that I'm a bit more skilled at optimising that framework. Anyway I'm not going to try and debunk things point-by-point, I'd rather avoid yet another unpleasant Julia flame-war.

by uoaei 4y ago

whimsicalism 4y ago

Because the `jax.numpy` operations & primitives are almost 1:1 with numpy, many working scientists who already have experience working with numpy will be able to figure out jax faster. It is also easier to rewrite existing code/snippets (say you were working on a non-differentiable simulator before) into jax if you already have them in numpy then to do the whole rewrite in pytorch. I will say that I think pytorch has improved its numpy compatability a lot in recent years, functions that I was convinced didn't exist with pytorch (like eigh) apparently actually do.

chazeon 4y ago

I have seen JAX-MD[1] but not sure about "much better". On the other hand, there is just no MD implemented with PyTorch. [1]: https://github.com/jax-md/jax-md https://github.com/jax-md/jax-md

new_user55 4y ago
There is [torch-md](https://github.com/torchmd/torchmd https://github.com/torchmd/torchmd)

patrickkidger 4y ago

At least prior to this announcement: JAX was much faster than PyTorch for differentiable physics. (Better JIT compiler; reduced Python-level overhead.) E.g for numerical ODE simulation, I've found that Diffrax (https://github.com/patrick-kidger/diffrax https://github.com/patrick-kidger/diffrax) is ~100 times faster than torchdiffeq on the forward pass. The backward pass is much closer, and for this Diffrax is about 1.5 times faster. It remains to be seen how PyTorch 2.0 will compare, of course! Right now my job is actually building out the scientific computing ecosystem in JAX, so feel free to ping me with any other questions.

adgjlsfhk1 4y ago
If you care about performance of differential physics you shouldn't use python. Diffrax is almost OKish, but is missing a ton of features (e.g. good stiff solvers, arbitrary precision support, events for anything other than stopping the simulation, ability to control the linear solve which are needed for large problems). For simple cases it can come close to the C++/Julia solvers, but for anything complicated, you either won't be able to formulate the model, or you won't be able to solve it efficiently.
- patrickkidger 4y ago
  > If you care about performance This definitely isn't true. On any benchmark I've tried, JAX and Julia basically match each other. Usually I find JAX to be a bit faster, but that might just be that I'm a bit more skilled at optimising that framework. Anyway I'm not going to try and debunk things point-by-point, I'd rather avoid yet another unpleasant Julia flame-war.