Automatic Differentiation with jax.grad

Introduction

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?

f(x) = x^2

Join the 1M+ learners on CodeSignal

Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal

Your First Gradient with jax.grad

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.

Python

import jax
import jax.numpy as jnp

def f(x):
    return x**3 + 2*x**2 - 5*x + 1

# Create a gradient function using jax.grad
grad_f = jax.grad(f)

# Evaluate at x = 2.0
x_val = 2.0
gradient_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 - 5
analytical_derivative_at_x = 3*(x_val**2) + 4*x_val - 5
print(f"Analytical derivative f'(x) at x={x_val}: {analytical_derivative_at_x}")

The above code outputs:

text

Value of x: 2.0
f(x) at x=2.0: 7.0
Gradient f'(x) at x=2.0 (using jax.grad): 15.0
Analytical 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 function grad_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) = x^3 + 2x^2 - 5x + 1$ , the derivative using calculus rules is $f'(x) = 3x^2 + 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.

Gradients of Multi-Variable Functions

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).

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 arguments
grad_g_wrt_ab = jax.grad(g, argnums=(0, 1))

a_val = 2.0
b_val = 3.0

print(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^2
dg_da_analytical = 2 * a_val * b_val
dg_db_analytical = a_val**2 + 3 * b_val**2

print(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}")

This snippet results in the following output:

text

Values: a=2.0, b=3.0
g(a,b): 39.0
Gradient dg/da (jax.grad): 12.0, Analytical: 12.0
Gradient dg/db (jax.grad): 31.0, Analytical: 31.0
Gradients (dg/da, dg/db) (jax.grad): (Array(12., dtype=float32, weak_type=True), Array(31., dtype=float32, weak_type=True))

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) = a^2 \cdot b + b^3$ evaluates to $39.0$ at $a = 2.0$ and $b = 3.0$ . The partial derivatives are $\frac{\partial g}{\partial a} = 2ab = 12.0$ and $\frac{\partial g}{\partial b} = a^2 + 3b^2 = 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.

Introduction

What is Automatic Differentiation?

Join the 1M+ learners on CodeSignal

Your First Gradient with jax.grad

Computing Values and Gradients Together

Gradients of Multi-Variable Functions

Conclusion and Next Steps