November 2024
Intermediate to advanced
408 pages
12h 7m
English
Overview
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:
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 Technology
Google’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 Book
Deep 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 Inside
About the Reader
For intermediate Python programmers who are familiar with deep learning.
About the Author
Grigory Sapunov holds a Ph.D. in artificial intelligence and is a Google Developer Expert in Machine Learning.
The technical editor on this book was Nicholas McGreivy.
Quotes
A comprehensive guide to mastering JAX, whether you’re a seasoned deep learning practitioner or just venturing into the realm of differentiable programming and large-scale numerical simulations.
- François Chollet, Software Engineer, Google
A must-read! The emphasis on functional programming has transformed the way I approach building models.
- Stephen Oates, Data Scientist, Allianz Australia
Great, modular code. Helpful explanations. This book is a treasure.
- Ritobrata Ghosh, Independent Deep Learning Consultant
I thoroughly enjoyed this excellent book! I feel confident that I can now apply JAX in my own work.
- Tony Holdroyd, Retired Senior Lecturer in Computer Science and Mathematics
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 calculations
- Build differentiable models with JAX primitives
- Run distributed and parallelized computations with JAX
- Use high-level neural network libraries such as Flax
- Leverage libraries and modules from the JAX ecosystem
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 Technology
Google’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 Book
Deep 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 Inside
- 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
About the Reader
For intermediate Python programmers who are familiar with deep learning.
About the Author
Grigory Sapunov holds a Ph.D. in artificial intelligence and is a Google Developer Expert in Machine Learning.
The technical editor on this book was Nicholas McGreivy.
Quotes
A comprehensive guide to mastering JAX, whether you’re a seasoned deep learning practitioner or just venturing into the realm of differentiable programming and large-scale numerical simulations.
- François Chollet, Software Engineer, Google
A must-read! The emphasis on functional programming has transformed the way I approach building models.
- Stephen Oates, Data Scientist, Allianz Australia
Great, modular code. Helpful explanations. This book is a treasure.
- Ritobrata Ghosh, Independent Deep Learning Consultant
I thoroughly enjoyed this excellent book! I feel confident that I can now apply JAX in my own work.
- Tony Holdroyd, Retired Senior Lecturer in Computer Science and Mathematics
Become an O’Reilly member and get unlimited access to this title plus top books and audiobooks from O’Reilly and nearly 200 top publishers, thousands of courses curated by job role, 150+ live events each month,
and much more.
Read now
Unlock full access
More than 5,000 organizations count on O’Reilly
O’Reilly covers everything we've got, with content to help us build a world-class technology community, upgrade the capabilities and competencies of our teams, and improve overall team performance as well as their engagement.Julian F.
I wanted to learn C and C++, but it didn't click for me until I picked up an O'Reilly book. When I went on the O’Reilly platform, I was astonished to find all the books there, plus live events and sandboxes so you could play around with the technology.Addison B.
I’ve been on the O’Reilly platform for more than eight years. I use a couple of learning platforms, but I'm on O'Reilly more than anybody else. When you're there, you start learning. I'm never disappointed.Amir M.
I'm always learning. So when I got on to O'Reilly, I was like a kid in a candy store. There are playlists. There are answers. There's on-demand training. It's worth its weight in gold, in terms of what it allows me to do.