Skip to content

← JAX Foundations step 18 of 25

Medium Primitives

Gradient of Packed Parameters

Jump to the task

Keep gradients aligned with parameters

A model can store its parameters in a dictionary instead of one long array. JAX calls nested containers of arrays pytrees. Differentiating a scalar loss with respect to a parameter dictionary returns a dictionary with the same keys and a gradient array for each parameter array.

import jax
import jax.numpy as jnp

params = {"w": jnp.array([1.0, 2.0]), "b": jnp.array([0.5])}
x = jnp.array([[1.0, -1.0]])

def loss(p, x):
    prediction = x @ p["w"] + p["b"]
    return jnp.sum(prediction ** 2)

grads = jax.grad(loss)(params, x)
# grads["w"] is [-1.0, 1.0]
# grads["b"] is [-1.0]

The prediction is -0.5. Squaring it gives a loss of 0.25. The derivative with respect to the prediction is -1, which is multiplied by each input feature to get the weight gradients. The bias gradient is also -1.

A single array is itself a pytree with one leaf. Its gradient is an array of the same shape, with no extra tuple or dictionary around it. This exercise uses that simpler representation: weights followed by a bias in one array. It does not ask you to flatten a nested dictionary.

Differentiate the parameter argument

argnums=0 selects the first positional argument, whether that argument is an array or a nested container. Keep the data argument fixed. Write the parameter slices inside the loss so JAX can follow their dependence on the array being differentiated.

The diagram below illustrates a separate nested parameter tree. Its numbered leaves are placeholders, not the packed input used here.

Loading visualization…

Problem

Implement grad_over_flat_params(flat_params, x) using jax.grad. Both inputs are real floating-point JAX arrays. flat_params has shape (d + 1,): the first d entries are weights and the last is the bias. The data array x has shape (N, d).

Compute predictions x @ weights + bias and a scalar loss equal to the sum of their squares. There are no target labels in this exercise. Return the gradient with respect to flat_params, in the same order and with shape (d + 1,). Use a sum, not a mean.

For flat_params=[1.0, 0.0, 0.0] and x=[[2.0, 3.0]], the prediction is 2. The two weight gradients are 2*2*2=8 and 2*2*3=12. The bias gradient is 2*2=4, so return [8.0, 12.0, 4.0].