Computing Gradients with `jax.grad`

1. The grad Transformation

In PyTorch, you are accustomed to an imperative, object-oriented approach to gradients: you define a computation, a scalar loss is produced, and you call loss.backward() to populate .grad attributes on tensors. JAX takes a different, more functional path. It provides a transformation, jax.grad, that takes a Python function and returns a new Python function that computes its gradient.

This aligns with the core philosophy of JAX: transforming functions. jax.grad operates on numerical functions that are:

• Composable from JAX-native functions.
• Return a single scalar float value.

Let's start with a video that introduces this concept and contrasts it with the imperative style you already know.

As the video shows, if you have a function f(x) that returns a scalar, jax.grad(f) returns a new function g(x) which computes ${\nabla }_{x}f(x)$.

A key strength of JAX's functional design is composability. Because jax.grad(f) is just another function, you can apply jax.grad to it again to compute higher-order derivatives. Given your mathematics background, you'll recognize the elegance of this.

import jax
import jax.numpy as jnp




# A simple polynomial
def f(x):
  return x**3 + 2*x**2 - 3*x + 1




# First derivative
grad_f = jax.grad(f)




# Second derivative
grad_f_2 = jax.grad(grad_f) 



# or more compactly, jax.grad(jax.grad(f))




# Third derivative
grad_f_3 = jax.grad(grad_f_2)

x_val = 2.0
print(f"f({x_val}) = {f(x_val)}")
print(f"f'({x_val}) = {grad_f(x_val)}")      # 3*x**2 + 4*x - 3 = 12 + 8 - 3 = 17
print(f"f''({x_val}) = {grad_f_2(x_val)}")     # 6*x + 4 = 12 + 4 = 16
print(f"f'''({x_val}) = {grad_f_3(x_val)}")    # 6

This ability to compose transformations is a recurring and powerful theme in JAX.

2. Differentiating With Respect to Specific Arguments: argnums

In any realistic machine learning model, your loss function will depend on more than just one input. Typically, it takes model parameters, input data, and labels.

def loss_fn(params, x, y): ...

By default, jax.grad computes the gradient with respect to the first argument (params in this case). But what if you need to differentiate with respect to another argument, or several at once? This is controlled by the argnums parameter.

• argnums=0 (the default): Differentiate with respect to the first argument.
• argnums=1: Differentiate with respect to the second argument.
• argnums=(0, 1): Differentiate with respect to both the first and second arguments. The resulting gradient function will return a tuple of gradients.

The following video segment provides a clear, hands-on demonstration of this.

This mechanism is fundamental for training models, as it allows you to precisely target the trainable parameters (often a pytree, which we'll cover soon) for gradient computation, while ignoring other inputs like data batches.

3. Handling Auxiliary Outputs: value_and_grad and has_aux

A common pattern for a loss function is to return not only the scalar loss value needed for differentiation but also other useful metrics for logging, like accuracy, precision, or internal model states.

def loss_and_metrics(params, batch):



    # ... computation ...
    loss = ...      # Scalar loss
    accuracy = ...  # Scalar accuracy
    return loss, accuracy

If you were to apply jax.grad to this function directly, you would get an error because the function does not return a single scalar. JAX needs to know which output is the loss to be differentiated.

This is handled by the has_aux flag. When you set has_aux=True, you are telling jax.grad that your function returns a pair: (value_to_differentiate, auxiliary_data). JAX will then compute the gradient of the first element and simply pass the second element through.

Furthermore, in a training step, you almost always need both the loss value (for logging) and the gradients (for updating weights). While you could call the function once to get the loss and the gradient function once to get the gradients, this is inefficient as it repeats the forward pass computation. JAX provides jax.value_and_grad as an efficient utility that computes both at the same time.

Let's watch a concise explanation of how has_aux works within this context.

4. A Practical Example: The Training Step

We can now combine these concepts—value_and_grad, argnums, and has_aux—to construct a typical training step function. We will use a well-structured example from a tutorial by the University of Amsterdam.

First, let's examine the loss function they define. It's a standard setup for binary classification that takes the model's state, parameters, and a data batch, returning both the loss and accuracy. This is a perfect candidate for our tools.

Now, with the calculate_loss_acc function in mind, let's see how jax.value_and_grad is used to create an efficient train_step. The following code block is the culmination of everything we've discussed.

Let's break down that critical grad_fn line from the resource:
grad_fn = jax.value_and_grad(calculate_loss_acc, argnums=1, has_aux=True)

• jax.value_and_grad: We are requesting both the function's output and its gradients.
• calculate_loss_acc: This is the function we're transforming.
• argnums=1: We specify that we want to compute gradients with respect to the second argument of calculate_loss_acc, which is params. This is where the model's trainable weights are.
• has_aux=True: We tell JAX that calculate_loss_acc returns a (loss, acc) pair. JAX will differentiate the loss and carry the acc along as auxiliary data.

Therefore, when grad_fn is called, it will return ((loss, acc), grads). The outer tuple separates the original function's output from the gradients. The inner tuple (loss, acc) is the original return value of calculate_loss_acc. This elegant, single line of code prepares everything needed for a complete training update.

Conclusion

In this lesson, you've moved beyond jit to master JAX's second core transformation, grad. You've seen that automatic differentiation in JAX is a functional transformation, providing a powerful and composable alternative to the imperative .backward() paradigm.

Key Takeaways:

• jax.grad(f) is a higher-order function that takes a scalar-output function f and returns its gradient function.
• Unlike the stateful updates in PyTorch, JAX's grad is a pure function transformation.
• The argnums argument allows you to precisely control which of a function's inputs are differentiated.
• The combination of jax.value_and_grad and has_aux=True is the idiomatic and efficient pattern for handling typical loss functions that return both a loss to be differentiated and auxiliary metrics to be passed through.

We have now covered compilation (jit) and differentiation (grad). In the next lesson, we will complete the trifecta of core JAX transformations by exploring jax.vmap, the tool for automatic vectorization. Mastering vmap is the final step before you can fluently compose these transformations to write clean, powerful, and highly performant code for batched deep learning.