Awesome JAX

for a more complete & current list see https://github.com/gn2cholas/awesome-jax

The JAX Library

• repo: https://github.com/google/jax
• docs: https://jax.readthedocs.io/en/latest/
• sysml'18 paper: https://research.google/pubs/pub47008/

Deep Learning libraries

• Flax: https://github.com/google/flax
• Trax: https://github.com/google/trax
• Haiku: https://github.com/google/trax See JAX repo & docs for a full list

Tools & Libraries

• 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

Research done with JAX

• 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

Popular Articles

• Massively parallel MCMC with JAX
• Meta-Learning in 50 Lines of JAX
• Normalizing Flows in 100 Lines of JAX
• Infinitely Wide Neural-Networks | Neural Tangents Explained
• Turbocharging SVD with JAX
• Creating Adversarial Examples for Neural Networks with JAX
• JAX: Differentiable Computing by Google
• Deep Learning with Jax and Elegy: Going beyond TensorFlow, Pytorch, and Keras
• Option Greeks in Python: JAX for automatic partial-differentiation of Black-Scholes

(More) Complete Lists of Reverse Dependencies

• wheelodex: https://www.wheelodex.org/projects/jax/rdepends/
• github: https://github.com/google/jax/network/dependents?package_id=UGFja2FnZS0zMjg2NDI1MjA%3D

Things that should be done with JAX ;)

• Analog cicuit simulator like SPICE
• APL/J/K/FP to XLA compiler