Deep Learning with JAX

Accelerate deep learning and other number-intensive tasks with JAX, Google’s awesome high-performance numerical computing library.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.In Deep Learning with JAX you will learn how to:Use JAX for numerical calculationsBuild differentiable models with JAX primitivesRun distributed and parallelized computations with JAXUse high-level neural network libraries such as FlaxLeverage libraries and modules from the JAX ecosystemDeep 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 technologyGoogle’s JAX offers a fresh vision for deep learning. This powerful library gives you fine control over low level processes like gradient calculations, delivering fast and efficient model training and inference, especially on large datasets. JAX has transformed how research scientists approach deep learning. Now boasting a robust ecosystem of tools and libraries, JAX makes evolutionary computations, federated learning, and other performance-sensitive tasks approachable for all types of applications.About the bookDeep Learning with JAX teaches you to build effective neural networks with JAX. In this example-rich book, you’ll discover how JAX’s unique features help you tackle important deep learning performance challenges, like distributing computations across a cluster of TPUs. You’ll put the library into action as you create an image classification tool, an image filter application, and other realistic projects. The nicely-annotated code listings demonstrate how JAX’s functional programming mindset improves composability and parallelization.What's insideUse JAX for numerical calculationsBuild differentiable models with JAX primitivesRun distributed and parallelized computations with JAXUse high-level neural network libraries such as FlaxAbout the readerFor intermediate Python programmers who are familiar with deep learning.About the authorGrigory Sapunov holds a Ph.D. in artificial intelligence and is a Google Developer Expert in Machine Learning. Table of ContentsPart 11 When and why to use JAX2 Your first program in JAXPart 23 Working with arrays4 Calculating gradients5 Compiling your code6 Vectorizing your code7 Parallelizing your computations8 Using tensor sharding9 Random numbers in JAX10 Working with pytreesPart 311 Higher-level neural network libraries12 Other members of the JAX ecosystemA Installing JAXB Using Google ColabC Using Google Cloud TPUsD Experimental parallelization