JAX-MD

DeepMind's neural network for ab-initio quantum chemistry, directly solving the many-electron Schrödinger equation via variational Monte Carlo with antisymmetric wavefunctions, extended to excited states (Phys. Rev. Research 2020, Science 2024)

Differentiable tokamak core transport simulator for fusion energy research, coupling PDE solvers with JAX auto-differentiation and neural-network surrogates for fast forward modelling, pulse-design, and trajectory optimization (Google DeepMind, Apache 2.0)

Machine learning toolkit for many-body quantum systems, implementing neural quantum states, variational Monte Carlo, and tensor network algorithms to solve ground-state and dynamical problems in condensed matter physics and quantum chemistry (EPFL & collaborators, Nature Physics 2019/2022+, 670+ stars)

First real quadrotor robot trained end-to-end with differentiable physics for vision-based agile flight, bridging simulation-based learning and real-world deployment with physics-informed neural network controllers (558+ stars)

Open-source JAX-based software suite for variational optimization of deep-learning molecular wave functions, solving electronic ground and excited states via neural-network trial wave functions with configurable FermiNet, PauliNet, Psiformer, LapNet, and DeepErwin ansätze, geometric transferability across molecular configurations, and effective core potential support (FU Berlin / Noé group, J. Chem. Phys. 2023, 420+ stars, MIT License)

Improved equivariant Transformer for 3D atomic graphs (ICLR2024)