Google JAX Cookbook
Google JAX Cookbook
Perform machine learning and numerical computing with combined capabilities of TensorFlow and NumPy
About the Book
This is the practical, solution-oriented book for every data scientists, machine learning engineers, and AI engineers to utilize the most of Google JAX for efficient and advanced machine learning. It covers essential tasks, troubleshooting scenarios, and optimization techniques to address common challenges encountered while working with JAX across machine learning and numerical computing projects.The book starts with the move from NumPy to JAX. It introduces the best ways to speed up computations, handle data types, generate random numbers, and perform in-place operations. It then shows you how to use profiling techniques to monitor computation time and device memory, helping you to optimize training and performance. The debugging section provides clear and effective strategies for resolving common runtime issues, including shape mismatches, NaNs, and control flow errors. The book goes on to show you how to master Pytrees for data manipulation, integrate external functions through the Foreign Function Interface (FFI), and utilize advanced serialization and type promotion techniques for stable computations.If you want to optimize training processes, this book has you covered. It includes recipes for efficient data loading, building custom neural networks, implementing mixed precision, and tracking experiments with Penzai. You'll learn how to visualize model performance and monitor metrics to assess training progress effectively. The recipes in this book tackle real-world scenarios and give users the power to fix issues and fine-tune models quickly.
Key Learnings
- Get your calculations done faster by moving from NumPy to JAX's optimized framework.
- Make your training pipelines more efficient by profiling how long things take and how much memory they use.
- Use debugging techniques to fix runtime issues like shape mismatches and numerical instability.
- Get to grips with Pytrees for managing complex, nested data structures across various machine learning tasks.
- Use JAX's Foreign Function Interface (FFI) to bring in external functions and give your computational capabilities a boost.
- Take advantage of mixed-precision training to speed up neural network computations without sacrificing model accuracy.
- Keep your experiments on track with Penzai. This lets you reproduce results and monitor key metrics.
- Use advanced visualization techniques, like confusion matrices and learning curves, to make model evaluation more effective.
- Create your own neural networks and optimizers directly in JAX so you have full control of the architecture.
- Use serialization techniques to save, load, and transfer models and training checkpoints efficiently.
Table of Content
- Transition NumPy to JAX
- Profiling Computation and Device Memory
- Debugging Runtime Values and Errors
- Mastering Pytrees for Data Structures
- Exporting and Serialization
- Type Promotion Semantics and Mixed Precision
- Integrating Foreign Functions (FFI)
- Training Neural Networks with JAX
Table of Contents
- Transition NumPy to JAX
- Profiling Computation and Device Memory
- Debugging Runtime Values and Errors
- Mastering Pytrees for Data Structures
- Exporting and Serialization
- Type Promotion Semantics and Mixed Precision
- Integrating Foreign Functions (FFI)
- Training Neural Networks with JAX
The Leanpub 60 Day 100% Happiness Guarantee
Within 60 days of purchase you can get a 100% refund on any Leanpub purchase, in two clicks.
Now, this is technically risky for us, since you'll have the book or course files either way. But we're so confident in our products and services, and in our authors and readers, that we're happy to offer a full money back guarantee for everything we sell.
You can only find out how good something is by trying it, and because of our 100% money back guarantee there's literally no risk to do so!
So, there's no reason not to click the Add to Cart button, is there?
See full terms...
Earn $8 on a $10 Purchase, and $16 on a $20 Purchase
We pay 80% royalties on purchases of $7.99 or more, and 80% royalties minus a 50 cent flat fee on purchases between $0.99 and $7.98. You earn $8 on a $10 sale, and $16 on a $20 sale. So, if we sell 5000 non-refunded copies of your book for $20, you'll earn $80,000.
(Yes, some authors have already earned much more than that on Leanpub.)
In fact, authors have earnedover $14 millionwriting, publishing and selling on Leanpub.
Learn more about writing on Leanpub
Free Updates. DRM Free.
If you buy a Leanpub book, you get free updates for as long as the author updates the book! Many authors use Leanpub to publish their books in-progress, while they are writing them. All readers get free updates, regardless of when they bought the book or how much they paid (including free).
Most Leanpub books are available in PDF (for computers) and EPUB (for phones, tablets and Kindle). The formats that a book includes are shown at the top right corner of this page.
Finally, Leanpub books don't have any DRM copy-protection nonsense, so you can easily read them on any supported device.
Learn more about Leanpub's ebook formats and where to read them