Welcome back to the third lesson of "JAX Fundamentals: NumPy Power-Up"! We're making fantastic progress on our journey to mastering JAX's most powerful capabilities. In our previous lessons, we explored JAX arrays and their immutable nature, then discovered how pure functions form the cornerstone of JAX's design. These concepts weren't just theoretical exercises — they were laying the groundwork for the truly transformative feature we'll explore today.
Today, we're diving into one of JAX's most celebrated features: automatic differentiation. This is where JAX begins to show its true power beyond just being a NumPy replacement. Automatic differentiation is the mathematical foundation that enables machine learning, optimization algorithms, and scientific computing applications to compute gradients effortlessly and accurately.
As you may recall from our previous lesson, pure functions are essential because they enable JAX's transformations to work reliably. Today, we'll see this principle in action as we use jax.grad to automatically compute derivatives of our pure functions. By the end of this lesson, you'll understand how to use jax.grad to compute gradients of scalar-output functions, evaluate these gradients at specific points, and even handle functions with multiple variables.
What is Automatic Differentiation?
Your First Gradient with jax.grad
Computing Values and Gradients Together
Gradients of Multi-Variable Functions
Conclusion and Next Steps
Congratulations! You've just mastered the fundamentals of automatic differentiation with jax.grad. You've learned how to transform any pure, scalar-output function into a gradient function using jax.grad, compute both function values and gradients efficiently with jax.value_and_grad, and handle multi-variable functions using the argnums parameter to specify which arguments to differentiate. Most importantly, you've seen how JAX transforms complex differentiation tasks that would be tedious by hand into simple, reliable computations that produce mathematically exact results.
This foundation in automatic differentiation opens the door to JAX's more advanced transformations and real-world applications. In our upcoming lessons, we'll explore how to combine jax.grad with other JAX transformations like jit for performance optimization, and we'll see how automatic differentiation forms the backbone of modern machine learning algorithms. The pure functions and gradient computations you've mastered today will be the building blocks for more sophisticated optimization and machine learning workflows ahead.
Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal
Before we start computing gradients with code, let's understand what automatic differentiation actually is and why it's so revolutionary for numerical computing. In calculus, we learned to compute derivatives by hand using rules like the power rule, product rule, and chain rule. For simple functions like f(x)=x2, finding the derivative f′(x)=2x is straightforward. But imagine trying to compute derivatives by hand for the complex functions found in modern machine learning models — functions with millions of parameters and hundreds of layers of computations!
Automatic differentiation (often abbreviated as autodiff) is a computational technique that allows computers to compute exact derivatives of functions defined by computer programs. Unlike numerical differentiation (which approximates derivatives using finite differences) or symbolic differentiation (which manipulates mathematical expressions), automatic differentiation computes exact derivatives by applying the chain rule systematically to the elementary operations in a program.
The key insight is that any complex function computed by a program can be broken down into a sequence of elementary operations like addition, multiplication, exponentials, and trigonometric functions. Each of these elementary operations has a known derivative, and automatic differentiation applies the chain rule to combine these derivatives automatically.
What makes this particularly powerful in JAX is that it works seamlessly with the pure functions we learned about previously. Since pure functions have no side effects and behave deterministically, JAX can safely analyze their computational structure and compute gradients without worrying about unpredictable behavior.
Now let's experience automatic differentiation in action with our first example. We'll start with a simple polynomial function and use jax.grad to compute its derivative automatically. This will demonstrate the basic workflow and show how JAX handles what we would normally do by hand in calculus.
import jaximport jax.numpy as jnpdef f(x): return x**3 + 2*x**2 - 5*x + 1# Create a gradient function using jax.gradgrad_f = jax.grad(f)# Evaluate at x = 2.0x_val = 2.0gradient_at_x = grad_f(x_val)function_value = f(x_val)print(f"Value of x: {x_val}")print(f"f(x) at x={x_val}: {function_value}")print(f"Gradient f'(x) at x={x_val} (using jax.grad): {gradient_at_x}")# Verify with analytical derivative# For f(x) = x^3 + 2x^2 - 5x + 1, the derivative is f'(x) = 3x^2 + 4x - 5analytical_derivative_at_x = 3*(x_val**2) + 4*x_val - 5print(f"Analytical derivative f'(x) at x={x_val}: {analytical_derivative_at_x}")
The above code outputs:
Value of x: 2.0f(x) at x=2.0: 7.0Gradient f'(x) at x=2.0 (using jax.grad): 15.0Analytical derivative f'(x) at x=2.0: 15.0
Let's break down what's happening here step by step. First, we define our function f(x) as a pure function that takes a single input and returns a scalar value. Notice that this function satisfies all the requirements for automatic differentiation: it's pure, deterministic, and returns a scalar output.
The magic happens with jax.grad(f). This doesn't immediately compute a gradient — instead, it returns a new functiongrad_f that computes the derivative of our original function. When we call grad_f(2.0), JAX automatically applies the differentiation rules and returns the value of the derivative at x=2.0.
Looking at the output, we see that for our function evaluated at x=2.0, the function value is 7.0 and the gradient is 15.0. To verify our result, we manually computed the analytical derivative: for f(x)=x3+2x2−5x+1, the derivative using calculus rules is f′(x)=3x2+4x−5. At x=2.0, this gives us f′(2.0)=3(4)+4(2)−5=15.0. The perfect match between JAX's automatic differentiation and our hand-calculated derivative demonstrates the accuracy of automatic differentiation.
In many practical applications, we need both the function's value and its gradient at the same point. Computing these separately would be inefficient, so JAX provides jax.value_and_grad to compute both simultaneously in a single, optimized pass through the computation:
Python
# Using value_and_grad for efficiencyvalue_and_grad_f = jax.value_and_grad(f)x_val_vgrad = 3.0value, gradient = value_and_grad_f(x_val_vgrad)print(f"Value of x: {x_val_vgrad}")print(f"f(x) (from value_and_grad): {value}")print(f"Gradient f'(x) (from value_and_grad): {gradient}")# Verify the value_and_grad resultanalytical_derivative_at_x_vgrad = 3*(x_val_vgrad**2) + 4*x_val_vgrad - 5print(f"Analytical derivative f'(x): {analytical_derivative_at_x_vgrad}")
Which outputs:
text
Value of x: 3.0f(x) (from value_and_grad): 31.0Gradient f'(x) (from value_and_grad): 34.0Analytical derivative f'(x): 34.0
The value_and_grad function is particularly important for optimization algorithms, where we typically need both the loss value (to monitor training progress) and the gradient (to update parameters). Instead of making two separate function calls, we get both pieces of information from a single, efficient computation.
Notice that value_and_grad_f(3.0) returns a tuple containing both the function value and its gradient at x=3.0. The output shows us that at x=3.0, the function evaluates to 31.0 with a gradient of 34.0. We verify this result analytically: f′(3.0)=3(9)+4(3)−5=34.0, confirming JAX's calculation once again.
The efficiency gain from value_and_grad comes from JAX's ability to reuse intermediate computations. When computing both the function value and its gradient, many of the intermediate calculations are shared between the forward pass (computing the value) and the backward pass (computing the gradient).
Real-world functions often depend on multiple variables, and we need to compute partial derivatives with respect to different inputs. JAX handles this elegantly through the argnums parameter, which specifies which arguments to differentiate with respect to. By default, jax.grad differentiates with respect to the first argument (argnums=0).
def g(a, b): return a**2 * b + b**3# Gradient with respect to first argument (default)grad_g_wrt_a = jax.grad(g, argnums=0)# Gradient with respect to second argument grad_g_wrt_b = jax.grad(g, argnums=1)# Gradients with respect to both argumentsgrad_g_wrt_ab = jax.grad(g, argnums=(0, 1))a_val = 2.0b_val = 3.0print(f"Values: a={a_val}, b={b_val}")print(f"g(a,b): {g(a_val, b_val)}")dg_da = grad_g_wrt_a(a_val, b_val)dg_db = grad_g_wrt_b(a_val, b_val)# Analytical partial derivatives for verification# dg/da = 2ab, dg/db = a^2 + 3b^2dg_da_analytical = 2 * a_val * b_valdg_db_analytical = a_val**2 + 3 * b_val**2print(f"Gradient dg/da (jax.grad): {dg_da}, Analytical: {dg_da_analytical}")print(f"Gradient dg/db (jax.grad): {dg_db}, Analytical: {dg_db_analytical}")grads_ab = grad_g_wrt_ab(a_val, b_val)print(f"Gradients (dg/da, dg/db) (jax.grad): {grads_ab}")
The argnums parameter gives us fine-grained control over which partial derivatives to compute. When argnums=0 (the default), jax.grad computes the partial derivative with respect to the first argument. When argnums=1, it computes the partial derivative with respect to the second argument.
Looking at our output, the function g(a,b)=a2⋅b+b3 evaluates to 39.0 at a=2.0 and b=3.0. The partial derivatives are ∂a∂g=2ab=12.0 and ∂b∂g=a2+3b2=31.0, which match JAX's computed gradients exactly.
Most powerfully, when we specify argnums=(0, 1), JAX returns a tuple of gradients — one for each specified argument. Notice in the output that the gradients are returned as JAX arrays with dtype=float32, which is JAX's default floating-point precision. The weak_type=True indicates that these arrays have weakly-typed semantics, meaning they can be promoted to higher precision if needed in computations.
This tuple format is incredibly useful for optimization algorithms that need to update multiple parameters simultaneously. The ability to compute all required partial derivatives in a single call makes JAX particularly efficient for machine learning applications.
Python
import jaximport jax.numpy as jnpdef f(x): return x**3 + 2*x**2 - 5*x + 1# Create a gradient function using jax.gradgrad_f = jax.grad(f)# Evaluate at x = 2.0x_val = 2.0gradient_at_x = grad_f(x_val)function_value = f(x_val)print(f"Value of x: {x_val}")print(f"f(x) at x={x_val}: {function_value}")print(f"Gradient f'(x) at x={x_val} (using jax.grad): {gradient_at_x}")# Verify with analytical derivative# For f(x) = x^3 + 2x^2 - 5x + 1, the derivative is f'(x) = 3x^2 + 4x - 5analytical_derivative_at_x = 3*(x_val**2) + 4*x_val - 5print(f"Analytical derivative f'(x) at x={x_val}: {analytical_derivative_at_x}")
text
Value of x: 2.0f(x) at x=2.0: 7.0Gradient f'(x) at x=2.0 (using jax.grad): 15.0Analytical derivative f'(x) at x=2.0: 15.0
Python
def g(a, b): return a**2 * b + b**3# Gradient with respect to first argument (default)grad_g_wrt_a = jax.grad(g, argnums=0)# Gradient with respect to second argument grad_g_wrt_b = jax.grad(g, argnums=1)# Gradients with respect to both argumentsgrad_g_wrt_ab = jax.grad(g, argnums=(0, 1))a_val = 2.0b_val = 3.0print(f"Values: a={a_val}, b={b_val}")print(f"g(a,b): {g(a_val, b_val)}")dg_da = grad_g_wrt_a(a_val, b_val)dg_db = grad_g_wrt_b(a_val, b_val)# Analytical partial derivatives for verification# dg/da = 2ab, dg/db = a^2 + 3b^2dg_da_analytical = 2 * a_val * b_valdg_db_analytical = a_val**2 + 3 * b_val**2print(f"Gradient dg/da (jax.grad): {dg_da}, Analytical: {dg_da_analytical}")print(f"Gradient dg/db (jax.grad): {dg_db}, Analytical: {dg_db_analytical}")grads_ab = grad_g_wrt_ab(a_val, b_val)print(f"Gradients (dg/da, dg/db) (jax.grad): {grads_ab}")
