Awesome JAX

The JAX Library

Neural Tangents
- code: https://github.com/google/neural-tangents
- paper: https://arxiv.org/abs/1912.02803
TensorNetwork
- code: https://github.com/google/TensorNetwork
JAX Molecular Dynamics
- code: https://github.com/google/jax-md
Alibaba Cloud Quantum Development Platform (ACQDP)
- code: https://github.com/alibaba/acqdp
Jraph - A library for graph neural networks in jax
- code: https://github.com/deepmind/jraph
SymJAX: symbolic CPU/GPU/TPU programming
- code: https://github.com/SymJAX/SymJAX
coax: Plug-n-play Reinforcement Learning in Python with OpenAI Gym and JAX
- cpde: https://github.com/microsoft/coax

Rethinking attention with Performers
- paper: https://arxiv.org/abs/2009.14794
Your Classifier is Secretly an Energy Based Model and You Should Treat it Like One
- paper: https://arxiv.org/abs/1912.03263
Lagrangean Neural Networks
- paper: https://arxiv.org/abs/2003.04630
Finite Versus Infinite Neural Networks: an Empirical Study
- paper: https://arxiv.org/abs/2007.15801
Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains
- paper: https://arxiv.org/abs/2006.10739
Reformer: The Efficient Transformer
- paper: https://arxiv.org/abs/2001.04451
Wide Neural Networks of Any Depth Evolve as Linear Models Under Gradient Descent
- paper: https://arxiv.org/abs/1902.06720