- Compiling machine learning programs via high-level tracing
- JAX, M.D. A Framework for Differentiable Physics
- Enabling Fast Differentially Private SGD via Just-in-Time Compilation and Vectorization
- JAX MD: End-to-End Differentiable, Hardware Accelerated, Molecular Dynamics in Pure Python
- Reimplementing Unirep in JAX