PyTorch 2.0

• jax-md

  2 1,085 7.4 Jupyter Notebook

  Differentiable, Hardware Accelerated, Molecular Dynamics

On the other hand, there is just no MD implemented with PyTorch.

[1]: https://github.com/jax-md/jax-md

• diffrax

  21 1,217 8.3 Python

  Numerical differential equation solvers in JAX. Autodifferentiable and GPU-capable. https://docs.kidger.site/diffrax/

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) 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, or 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.

• WorkOS

  workos.com sponsored

  The modern identity platform for B2B SaaS. The APIs are flexible and easy-to-use, supporting authentication, user identity, and complex enterprise features like SSO and SCIM provisioning.

  

• tinygrad

  58 17,800 9.7 Python

  Discontinued You like pytorch? You like micrograd? You love tinygrad! ❤️ [Moved to: https://github.com/tinygrad/tinygrad] (by geohot)

  

• torchmd

  1 505 7.9 Python

  End-To-End Molecular Dynamics (MD) Engine using PyTorch

  

Suggest a related project