Part 3. Ecosystem

Part 3 introduces you to the vibrant ecosystem surrounding JAX, showcasing the libraries and tools that extend its functionality into various domains of deep learning and beyond. This section highlights how JAX fits into a broader context, enabling you to leverage its power in conjunction with other specialized libraries. Through two comprehensive chapters, you’ll discover the high-level neural network libraries that simplify model building and training, as well as other members of the JAX ecosystem that cater to a wide array of scientific and computational needs.

Chapter 11 discusses higher-level neural network libraries, focusing on Flax and its Linen API, along with Optax for gradient transformations. You’ll learn how to construct models more intuitively, manage training states, and interact with the Hugging Face ecosystem to access state-of-the-art models. This chapter bridges the gap between JAX’s core capabilities and the practical needs of deep learning projects, providing you with the tools to build, train, and deploy complex models efficiently.

Chapter 12 takes a broader look at the JAX ecosystem, showcasing libraries for various machine learning tasks, including reinforcement learning and evolutionary computations. We’ll also explore JAX modules for other scientific fields like physics, chemistry, and more.