Introduction to JAX Ecosystem

At this point in the course, JAX doesn’t seem much like a specialized neural network library, despite its wide suite of features. This is understandable given JAX, at its core, is a parallelizable NumPy with a number of other features.

The story doesn’t conclude there, though. In order to facilitate deep learning, JAX has its own ecosystem of libraries, comprising of:

1. Haiku
2. Optax
3. Jraph
4. RLax
5. Chex

Each of these libraries deserves a chapter of its own, and even that would be unable to fully unpack their particular details. However, we’ll just provide a brief overview of each in this introductory course.

Note: Flax is a library for neural networks but is usually not considered part of the JAX ecosystem.