# 🔗 Chain rule for composed functions

By Better Wise

## Why compositions need their own rule

You already have the two ingredients: the gradient vector of a scalar function, and the Jacobian of a vector-valued function. Almost nothing you differentiate in practice arrives expanded, because most functions are compositions: a linear map feeding an elementwise activation, a log wrapping an exponential, a sum of squares wrapping an affine map. The chain rule is the one rule that handles all of them without expanding anything.

After this lesson you can:

- name the intermediate variables of a composition and write the forward pass;
- assemble the Jacobian product in the correct order;
- collapse that product to a gradient when the outer map is scalar;
- confirm every hand-derived gradient with a central finite difference in Python.

One-line recall from the previous two lessons: the gradient of $f:\mathbb{R}^n \to \mathbb{R}$ is the column vector of partials; the Jacobian of $h:\mathbb{R}^n \to \mathbb{R}^m$ is the $m \times n$ matrix whose row $i$ is $\nabla h_i$.

## The chain rule in Jacobian form

Write the composition as a chain of named maps, $f = g \circ h$, so $x \mapsto h(x) = y \mapsto g(y) = f(x)$. With $h:\mathbb{R}^n \to \mathbb{R}^m$ and $g:\mathbb{R}^m \to \mathbb{R}^k$, the derivative of the composition is the product of derivatives evaluated at matching points:

$$
J_f(x) = J_g(h(x))\, J_h(x).
$$

Shapes confirm the order: $(k \times m)(m \times n) = k \times n$. The inner Jacobian is evaluated at $x$; the outer Jacobian is evaluated at _____. Reversing the order gives a shape mismatch, which is why the ordering mistake is usually caught by arithmetic rather than by luck.

```mermaid
flowchart TD
    X["🧮<br/>x"] --> H["🔁<br/>h(x) = y"]
    H --> G["🎯<br/>g(y) = f(x)"]
    JX["J_h(x)"] --> P["J_g(h(x)) J_h(x)"]
    JY["J_g(h(x))"] --> P
```

When $k = 1$, $J_g$ is a single row, the row vector $\nabla g(y)^\top$, and the product becomes a row vector. Transposing gives the gradient of the composition, which is worth memorizing in its own right:

$$
\nabla f(x) = J_h(x)^\top \nabla g(h(x)).
$$

In this form the cloze reads: the gradient of a scalar-valued composition is _____, or the inner Jacobian transposed times the outer gradient.

## Worked example: a nested scalar composition

Let $f(x) = e^{x_1^2 + 3x_2}$, so $h(x) = x_1^2 + 3x_2$ and $g(u) = e^u$. Here $J_h(x) = [2x_1,\; 3]$ and $g'(u) = e^u$, so

$$
\nabla f(x) = e^{x_1^2+3x_2}\begin{bmatrix} 2x_1 \\ 3 \end{bmatrix}.
$$

Now wrap a log around it: $L = \log f$. Then $L = x_1^2 + 3x_2$ and $\nabla L = [2x_1, 3]^\top$. The scalar factors $f$ and $1/f$ cancel, which is the chain rule telling you composition can simplify, not complicate.

## Worked example: composing through a matrix-shaped Jacobian

Take $h:\mathbb{R}^2 \to \mathbb{R}^2$ with $h(x) = (x_1 x_2,\; x_1 + x_2^2)$ and $g(y) = y_1 y_2$.

Hand derivation:

$$
J_h(x) = \begin{bmatrix} x_2 & x_1 \\ 1 & 2x_2 \end{bmatrix},
\qquad \nabla g(y) = \begin{bmatrix} y_2 \\ y_1 \end{bmatrix}.
$$

Multiplying $J_h(x)^\top \nabla g(y)$ and substituting $y = h(x)$:

$$
\nabla f(x) = \begin{bmatrix} x_2 y_2 + y_1 \\ x_1 y_2 + 2x_2 y_1 \end{bmatrix} = \begin{bmatrix} 2x_1x_2 + x_2^3 \\ x_1^2 + 3x_1x_2^2 \end{bmatrix}.
$$

Direct expansion gives $f = x_1^2 x_2 + x_1 x_2^3$, and differentiating that term by term lands on exactly the same vector. Two routes, one answer.

## Worked example: elementwise layers produce diagonal Jacobians

Compositions that apply a scalar function to each coordinate are everywhere. For $a = Wx + b$ with $W \in \mathbb{R}^{2 \times 2}$, $t = \tanh(a)$ (elementwise), and readout $f = r^\top t$:

$$
J_a(x) = W, \qquad J_t(a) = \operatorname{diag}(1 - t_1^2,\, 1 - t_2^2), \qquad \nabla_t f = r.
$$

Chaining: $\nabla_x f = W^\top \operatorname{diag}(1-t^2)\, r = W^\top\big((1 - t^2) \odot r\big)$. An elementwise map always contributes a _____ Jacobian, so the matrix product degrades into an elementwise scaling. That is the whole reason backpropagation is cheap.

## Confirming the chain with finite differences

Never ship a hand-derived gradient without a numeric check. The central difference for coordinate $i$ is _____, which is second-order accurate: its truncation error is $O(h^2)$, while a one-sided difference is only $O(h)$.

```python
import numpy as np

def f(x):                                       # g o h, R^2 -> R
    y1, y2 = x[0] * x[1], x[0] + x[1]**2
    return y1 * y2

def grad_analytic(x):
    x1, x2 = x
    return np.array([2*x1*x2 + x2**3, x1**2 + 3*x1*x2**2])

def grad_fd(fun, x, h=1e-6):
    g = np.zeros_like(x)
    for i in range(x.size):
        e = np.zeros_like(x); e[i] = h
        g[i] = (fun(x + e) - fun(x - e)) / (2 * h)
    return g

x = np.array([1.3, -0.7])
ga, gn = grad_analytic(x), grad_fd(f, x)
print(ga, gn, np.linalg.norm(ga - gn) / np.linalg.norm(gn))
```

The relative error comes out around $10^{-10}$. Anything below $10^{-6}$ is a pass; a relative error near $10^{-2}$ means a missing factor, a wrong transpose, or the outer Jacobian evaluated at $x$ instead of at $h(x)$.

## Practical rules for chaining derivatives

- **Name your intermediates.** Write $a = Wx + b$, $t = \tanh(a)$, $f = r^\top t$ before touching a matrix. Most chain-rule errors are bookkeeping errors.
- **Check shapes first.** If a product does not have the shape of the answer, the order is wrong. Do this before checking values.
- **Mind the evaluation points.** Inner Jacobian at the input, outer Jacobian at the intermediate value. This is the single most common bug in hand derivations.
- **Elementwise means diagonal.** Never build the diagonal explicitly in code; write it as an elementwise product.
- **Use $h \approx 10^{-6}$ for central differences in float64.** Smaller $h$ lets roundoff dominate; larger $h$ lets truncation dominate.
- **Test at a random point, not a tidy one.** Zeros and ones hide sign and factor errors.

## What to carry forward

The chain rule says the derivative of a composition is the ordered product of the derivatives of its parts, with the gradient case being $J_h(x)^\top \nabla g(h(x))$; everything else in this lesson is bookkeeping around that one line, checked numerically.

Next comes the checkpoint on first-order derivatives, where you will be asked to move between gradients, Jacobians, and chained compositions on demand.

Front: Jacobian of a composition f = g ∘ h
Back: J_f(x) = J_g(h(x)) J_h(x), outer times inner.

Front: Gradient of a scalar-valued composition f = g ∘ h
Back: ∇f(x) = J_h(x)^T ∇g(h(x)).

Front: At which point is the outer Jacobian J_g evaluated?
Back: At the intermediate value h(x), not at x.

Front: At which point is the inner Jacobian J_h evaluated?
Back: At the input x.

Front: Jacobian of an elementwise function t = φ(a)
Back: A diagonal matrix diag(φ'(a_i)).

Front: Jacobian of the affine map a = Wx + b
Back: Just W, independent of x.

Front: Gradient of f(x) = e^(x1² + 3x2)
Back: e^(x1² + 3x2) times [2x1, 3].

Front: Central difference formula for coordinate i
Back: (f(x + h e_i) - f(x - h e_i)) / (2h).

Front: Truncation error order of a central difference
Back: O(h²), versus O(h) for a one-sided difference.

Front: Why not use h = 1e-14 in a finite-difference check?
Back: Roundoff in the subtraction swamps the difference; 1e-6 is a good float64 default.

Front: Gradient of f = r^T tanh(Wx + b)
Back: W^T ((1 - tanh(Wx+b)²) ⊙ r).

Front: Shape test that catches reversed chain-rule products
Back: (k×m)(m×n) must equal the k×n Jacobian of the composition.

Front: What does a relative error near 1e-2 in a finite-difference check usually mean?
Back: A missing factor, a wrong transpose, or the outer Jacobian evaluated at the wrong point.
A last recall sweep before the practical rules.

The multivariate chain rule for Jacobians reads J_f(x) = _____, outer times inner.
For a scalar-valued composition f = g ∘ h the gradient is ∇f(x) = _____.
An elementwise nonlinearity t = φ(a) has a _____ Jacobian, diag(φ'(a_i)).
A central difference has truncation error _____, one order better than a one-sided difference.
A reliable default step for a float64 finite-difference check is _____.

## Rules of thumb that save you from silent gradient bugs

- **Write the shapes before the algebra.** Label every intermediate with its dimension, for example x is n, a is m, f is a scalar. If the product you wrote cannot be formed, the chain rule is written in the wrong order, not the wrong function.
- **The outer Jacobian never sees x.** A very common slip is writing J_g(x) J_h(x) for f = g ∘ h. The outer factor lives at the intermediate value h(x). Numerically this produces a plausible-looking but wrong gradient, which is why the finite-difference check matters.
- **Transpose only when you leave Jacobian land.** Keep everything as Jacobians while chaining, then apply the transpose once, at the end, if you want a gradient rather than a row vector. Transposing mid-chain is the second most common error.
- **Do not shrink h out of enthusiasm.** h = 1e-12 in float64 gives garbage because the two function values agree to twelve digits and the subtraction loses them. Start at 1e-6, and if the error looks large, try 1e-5 and 1e-7 to see whether roundoff or truncation dominates.
- **Judge with relative error, not absolute error.** Compare with `np.linalg.norm(num - ana) / (np.linalg.norm(num) + np.linalg.norm(ana))`. Values near 1e-7 mean agreement; values near 1e-2 mean a missing factor, a wrong transpose, or the outer Jacobian at the wrong point.
- **Reuse one checker.** A single function that takes a callable and its analytic gradient and loops over coordinates will catch more bugs over the course than any amount of staring at algebra.

```python
import numpy as np

def check_grad(f, grad_f, x, h=1e-6):
    x = np.asarray(x, dtype=float)
    num = np.zeros_like(x)
    for i in range(x.size):
        e = np.zeros_like(x); e[i] = h
        num[i] = (f(x + e) - f(x - e)) / (2 * h)
    ana = np.asarray(grad_f(x), dtype=float)
    rel = np.linalg.norm(num - ana) / (np.linalg.norm(num) + np.linalg.norm(ana) + 1e-12)
    return rel, num, ana
```

## What to carry forward

You now have the full first-order toolkit: partials assemble into a gradient, vector-valued maps get a Jacobian, and compositions multiply those Jacobians from the outside in, with the outer factor evaluated at the intermediate value. The reusable habit is the pairing: derive by hand, then confirm with a central difference at 1e-6 and a relative error near 1e-7.

The next lesson is the module checkpoint, where you will derive a few gradients and Jacobians from scratch, chain them, and defend each answer with a numeric check, with no new theory introduced.

Front: Chain rule for f = g ∘ h, in Jacobian form
Back: J_f(x) = J_g(h(x)) J_h(x), outer times inner.

Front: Chain rule for a scalar-valued composition, in gradient form
Back: ∇f(x) = J_h(x)^T ∇g(h(x)).

Front: Where is the outer Jacobian evaluated?
Back: At h(x), the intermediate value, never at x.

Front: Jacobian of an elementwise map t = φ(a)
Back: The diagonal matrix diag(φ'(a_i)).

Front: Jacobian of the affine map a = Wx + b
Back: W, and it does not depend on x.

Front: Gradient of f(x) = exp(x1² + 3x2)
Back: exp(x1² + 3x2) · [2x1, 3].

Front: Gradient of f = rᵀ tanh(Wx + b)
Back: Wᵀ ((1 − tanh(Wx+b)²) ⊙ r).

Front: Central difference for coordinate i
Back: (f(x + h e_i) − f(x − h e_i)) / (2h).

Front: Truncation order of central vs one-sided differences
Back: O(h²) for central, O(h) for one-sided.

Front: Why avoid h = 1e-14 in float64?
Back: Roundoff in the subtraction swamps the difference; 1e-6 is a sound default.

Front: Shape test for a chained product
Back: (k×m)(m×n) must give the k×n Jacobian of the composition.

Front: Relative error near 1e-2 in a gradient check usually means
Back: A missing factor, a wrong transpose, or the outer Jacobian evaluated at the wrong point.

---

## Contents

Part of [🧮 Mathematics for Machine Learning 2: Calculus and Optimization for Training Models](https://everwise.study/courses/mathematics-for-machine-learning-2-calculus-and-optimization-for-training.md)

- **📐 First-order derivatives: gradients, Jacobians and the chain rule**
  - [📐 From partial derivatives to the gradient vector](https://everwise.study/courses/mathematics-for-machine-learning-2-calculus-and-optimization-for-training/partial-derivatives-to-gradient-vector.md)
  - [🧮 Jacobians of vector-valued functions](https://everwise.study/courses/mathematics-for-machine-learning-2-calculus-and-optimization-for-training/jacobian-vector-valued-functions.md) (sign-in required)
  - [🔗 Chain rule for composed functions](https://everwise.study/courses/mathematics-for-machine-learning-2-calculus-and-optimization-for-training/chain-rule-composed-functions.md) — this lesson
  - [✅ Checkpoint: first-order derivatives](https://everwise.study/courses/mathematics-for-machine-learning-2-calculus-and-optimization-for-training/checkpoint-first-order-derivatives.md)
- **📉 Second-order behavior and convexity**
  - [🌀 Hessians and second-order behavior](https://everwise.study/courses/mathematics-for-machine-learning-2-calculus-and-optimization-for-training/hessians-second-order-behavior.md)
  - [📈 Taylor expansions, convexity and local minima](https://everwise.study/courses/mathematics-for-machine-learning-2-calculus-and-optimization-for-training/taylor-expansions-convexity-local-minima.md)
  - [✅ Checkpoint: curvature and convexity](https://everwise.study/courses/mathematics-for-machine-learning-2-calculus-and-optimization-for-training/checkpoint-curvature-and-convexity.md)
- **⛰️ Constrained optimization and gradient descent**
  - [⚖️ Lagrange multipliers for constrained problems](https://everwise.study/courses/mathematics-for-machine-learning-2-calculus-and-optimization-for-training/lagrange-multipliers-constrained-problems.md)
  - [🔽 Gradient descent and learning rate behavior](https://everwise.study/courses/mathematics-for-machine-learning-2-calculus-and-optimization-for-training/gradient-descent-learning-rate-behavior.md) (premium content)
  - [✅ Checkpoint: constrained optimization and descent](https://everwise.study/courses/mathematics-for-machine-learning-2-calculus-and-optimization-for-training/checkpoint-constrained-optimization-descent.md) (premium content)
- [Course cheatsheet](https://everwise.study/courses/mathematics-for-machine-learning-2-calculus-and-optimization-for-training/cheatsheet.md) (premium content)