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

1. Transition NumPy to JAX
2. Profiling Computation and Device Memory
3. Debugging Runtime Values and Errors
4. Mastering Pytrees for Data Structures
5. Exporting and Serialization
6. Type Promotion Semantics and Mixed Precision
7. Integrating Foreign Functions (FFI)
8. Training Neural Networks with JAX

"Learning PyTorch 2.0, Second Edition" is a fast-learning, hands-on book that emphasizes practical PyTorch scripting and efficient model development using PyTorch 2.3 and CUDA 12. This edition is centered on practical applications and presents a concise methodology for attaining proficiency in the most recent features of PyTorch.

The book presents a practical program based on the fish dataset which provides step-by-step guidance through the processes of building, training and deploying neural networks, with each example prepared for immediate implementation. Given your familiarity with machine learning and neural networks, this book offers concise explanations of foundational topics, allowing you to proceed directly to the practical, advanced aspects of PyTorch programming. The key learnings include the design of various types of neural networks, the use of torch.compile() for performance optimization, the deployment of models using TorchServe, and the implementation of quantization for efficient inference.

Furthermore, you will also learn to migrate TensorFlow models to PyTorch using the ONNX format. The book employs essential libraries, including torchvision, torchserve, tf2onnx, onnxruntime, and requests, to facilitate seamless integration of PyTorch with production environments. Regardless of whether the objective is to fine-tune models or to deploy them on a large scale, this second edition is designed to ensure maximum efficiency and speed, with practical PyTorch scripting at the forefront of each chapter.

Key Learnings

• Master tensor manipulations and advanced operations using PyTorch's efficient tensor libraries.
• Build feedforward, convolutional, and recurrent neural networks from scratch.
• Implement transformer models for modern natural language processing tasks.
• Use CUDA 12 and mixed precision training (AMP) to accelerate model training and inference.
• Deploy PyTorch models in production using TorchServe, including multi-model serving and versioning.
• Migrate TensorFlow models to PyTorch using ONNX format for seamless cross-framework compatibility.
• Optimize neural network architectures using torch.compile() for improved speed and efficiency.
• Utilize PyTorch's Quantization API to reduce model size and speed up inference.
• Setup custom layers and architectures for neural networks to tackle domain-specific problems. 
• Monitor and log model performance in real-time using TorchServe's built-in tools and configurations.

Table of Content

1. Introduction To PyTorch 2.3 and CUDA 12
2. Getting Started with Tensors
3. Building Neural Networks with PyTorch
4. Training Neural Networks
5. Advanced Neural Network Architectures
6. Quantization and Model Optimization
7. Migrating TensorFlow to PyTorch
8. Deploying PyTorch Models with TorchServe