In Deep Learning with JAX you will learn how to
β’ Use JAX for numerical calculations
β’ Build differentiable models with JAX primitives
β’ Run distributed and parallelized computations with JAX
β’ Use high-level neural network libraries such as Flax and Haiku
β’ Leverage libraries and modules from the JAX ecosystem
The JAX numerical computing library tackles the core performance challenges at the heart of deep learning and other scientific computing tasks. By combining Googleβs Accelerated Linear Algebra platform (XLA) with a hyper-optimized version of NumPy and a variety of other high-performance features, JAX delivers a huge performance boost in low-level computations and transformations.
Deep Learning with JAX is a hands-on guide to using JAX for deep learning and other mathematically-intensive applications. Google Developer Expert Grigory Sapunov steadily builds your understanding of JAXβs concepts. The engaging examples introduce the fundamental concepts on which JAX relies and then show you how to apply them to real-world tasks. Youβll learn how to use JAXβs ecosystem of high-level libraries and modules, and also how to combine TensorFlow and PyTorch with JAX for data loading and deployment.
About the book
Deep Learning with JAX teaches you how to use JAX and its ecosystem to build neural networks. Youβll learn by exploring interesting examples including an image classification tool, an image filter application, and a massive scale neural network with distributed training across a cluster of TPUs. Discover how to work with JAX for hardware and other low-level aspects and how to solve common machine learning problems with JAX. By the time youβre finished with this awesome book, youβll be ready to start applying JAX to your own research and prototyping.
About the reader
For intermediate Python programmers who are famil