Many computational tasks, especially in machine learning model training, rely on calculating the gradient of functions. Automatic differentiation provides an efficient way to compute these derivatives. This chapter introduces jax.grad, the core JAX transformation for obtaining gradient functions from Python code operating on numerical inputs.

You will learn how to:

• Apply jax.grad to compute the gradient $∇f(x)$ of a scalar-valued function $f$.
• Understand the basics of reverse-mode automatic differentiation as used by grad.
• Control differentiation with respect to specific function arguments.
• Compute higher-order derivatives by composing grad.
• Efficiently get both the function's output value and its gradient using jax.value_and_grad.
• Recognize how control flow interacts with differentiation and identify potential limitations.

By the end of this chapter, you will be able to use jax.grad effectively to differentiate your numerical functions within the JAX framework.