JAX
JAX is a Python library for high-performance array-based numerical computing with composable transformations including automatic differentiation, compilation, vectorization, and parallelization.
Deep Learning with JAX
A functional-first numerical computing library for machine learning and scientific computing in which state is made explicit and transformations like autodiff and compilation are central.
- How it is used
- Deploys NumPy-style numerical programs for differentiable, compiled, vectorized, and parallel machine-learning or scientific computation.
- How deep it goes
- Explains arrays, functional state, four core transformations, hardware, and ecosystem tradeoffs, while deferring detailed chapters on each feature.
- Examples
- Uses immutable indexed updates, performance comparisons, and an MNIST neural network combining gradients, compilation, and auto-vectorization.
- What you can do after
- Build a simple MNIST neural network using explicit state and grad(), jit(), and vmap().
Deep Learning with R, Third Edition
JAX is a Python library for high-performance numerical computing and machine learning with composable transformations such as automatic differentiation and just-in-time compilation.
- How it is used
- Positions JAX as a low-level tensor and differentiation engine for neural-network training, either directly or beneath Keras.
- How deep it goes
- Develops a full stateless linear-classifier loop and a Keras custom training step, without deriving JAX internals.
- Examples
- Works through a linear classifier and an MNIST Keras training step using explicit state and gradients.
- What you can do after
- Build it from scratch: implement a stateless linear-classifier training loop with `value_and_grad()` and `jit()`.
Deep Learning with Python, Third Edition
A numerical computing and machine learning library for high-performance array computation, automatic differentiation, and compilation.
- How it is used
- Supports low-level differentiable computation, custom training loops, Keras backend execution, and gradient-based model interpretation.
- How deep it goes
- Works through a stateless linear-classifier loop and a more involved Keras training step, including gradients, updated state, and compilation.
- Examples
- Uses a linear classifier, a custom Keras training step, image gradient ascent, and predicted-class gradients.
- What you can do after
- Build it from scratch: implement a basic linear classifier and its stateless JAX training loop.
Read how Deep Learning with Python, Third Edition explains JAX →
Related concepts
- XLAuses · XLA is a compiler that optimizes and compiles linear-algebra and array computations for supported target hardware.
- PyTorchcontrasts with · PyTorch is an open-source machine-learning framework for Python that provides tensor computation, automatic differentiation, and tools for building and training neural networks.
- TPUuses · A Tensor Processing Unit (TPU) is a specialized processor designed to accelerate tensor computations, particularly for machine-learning workloads.
- automatic differentiationuses · Automatic differentiation is a computational technique that evaluates derivatives of a program by systematically applying the chain rule to its elementary operations.
- jax.numpyimplements · jax.numpy is JAX’s NumPy-compatible API for creating and manipulating arrays while supporting JAX transformations.
- backendkind of · A backend is a software component or subsystem that provides underlying processing, data access, execution, or services to clients or other components, typically behind a user-facing interface.
These guides are generated from the text of Manning books by Manning's knowledge graph and link to the sections they draw on. The full books are available on liveBook and at manning.com.