Jax.grad
jax.grad is a JAX transformation that creates a function computing derivatives, typically gradients of a scalar-output function with respect to specified inputs.
Deep Learning with JAX
A JAX transformation that returns a derivative-computing version of a function, typically with respect to its first argument by default.
- How it is used
- Calculates gradients for loss functions, parameter updates, autodiff experiments, and composable numerical programs.
- How deep it goes
- Explains functional gradients, argnums, auxiliary outputs, nested parameters, stopping gradient flow, and higher derivatives, while deferring autodiff-mode internals.
- Examples
- Demonstrates linear regression, MNIST loss gradients, Minkowski distance, dictionary parameters, higher derivatives, and per-sample gradients.
- What you can do after
- Apply it via JAX to obtain parameter, selected-argument, nested-structure, higher-order, or per-sample gradients.
Deep Learning with Python, Third Edition
A JAX function that creates a gradient-computing function for differentiable inputs and outputs.
- How it is used
- Computes gradients for loss functions used in array programs, model training, and input-image optimization.
- How deep it goes
- Explains function transformation, scalar-loss constraints, argument conventions, auxiliary outputs, and a worked example without deriving automatic differentiation internals.
- Examples
- Uses a squared scalar loss, structured model state, auxiliary outputs, and stochastic gradient ascent on an input image.
- What you can do after
- Apply it via JAX to create reusable gradient functions for scalar losses and use them in optimization steps.
Read how Deep Learning with Python, Third Edition explains jax.grad →
Related concepts
- gradientproduces · A gradient is the vector of partial derivatives of a scalar function with respect to its variables, indicating its local rate and direction of change.
- JAXused by · JAX is a Python library for high-performance array-based numerical computing with composable transformations including automatic differentiation, compilation, vectorization, and parallelization.
- jax.value_and_gradalternative to · jax.value_and_grad is a JAX transformation that creates a function returning both a scalar function’s value and its gradient.
- automatic differentiationimplements · Automatic differentiation is a computational technique that evaluates derivatives of a program by systematically applying the chain rule to its elementary operations.
- vectorizationuses · Vectorization is the representation of data or computations as vectors or arrays so they can be processed numerically or collectively.
- objective functionuses · An objective function assigns values to candidate solutions, predictions, or parameter settings according to a criterion to be minimized or maximized.
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.