Introduction to JAX Ecosystem

This chapter provides an overview of the 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

Get hands-on with 1000+ tech skills courses.