4 ms·
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
by 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 agoIf 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.