The JAX Ecosystem is a thriving community of JAX libraries. Here are some libraries developed at DeepMind:
| Topic | Library |
|---|---|
| Neural networks | Haiku |
| Optimizers | Optax |
| Probability distributions | Distrax |
| Image processing | PIX |
| GraphNets | Jraph |
| Protein folding | AlphaFold |
| Reinforcement learning | RLax |
| AlphaZero-like Tree Search | Mctx |
| Testing & debugging | Chex |
| Second-Order Optimization | KFAC-JAX |
| Mixed Precision Training | JMP |
| Audio processing | AUX |
| TensorFlow to JAX conversion | tf2jax |
Originally, DeepMind introduced the JAX Ecosystem in a blog post.