Awesome JAX Libraries

We have just seen the tip of the iceberg here in the course when it comes to JAX libraries.
Here is a curated list of awesome JAX libraries, taken from Awesome JAX.

---

Libraries

• Neural Network Libraries
  • Flax: Centered on flexibility and clarity.
  • Haiku: Focused on simplicity, created by the authors of Sonnet at DeepMind.
  • Objax: Has an object oriented design similar to PyTorch.
  • Elegy: A high-level API for deep learning in JAX. Supports Flax, Haiku, and Optax.
  • Trax: “Batteries included” deep learning library focused on providing solutions for common workloads.
  • Jraph: Lightweight graph neural network library.
  • Neural Tangents: A high-level API for specifying neural networks of both finite and infinite width.

...