CALCULUS / 7. MULTIVARIABLE CHAIN RULE

Multivariable Chain Rule & Jacobian

Gradients for vector functions — the full generalization

EXPLANATION

The multivariable chain rule extends the chain rule to functions with multiple inputs and outputs.

For f(g(x), h(x)) — one output, two intermediate functions:
df/dx = (∂f/∂g)(dg/dx) + (∂f/∂h)(dh/dx)

In vector form (from Raschka's appendix):
df/dx = ∇f · v'(x)  where v = [g(x), h(x)]ᵀ

The Jacobian matrix: when f: Rⁿ → Rᵐ (vector input, vector output):
J[i,j] = ∂fᵢ/∂xⱼ

The Jacobian is an m×n matrix containing all first-order partial derivatives.

For a single function f: Rⁿ → R (like a loss function):
• Gradient ∇f is a column vector (n×1)
• Jacobian Jf is a row vector (1×n) = ∇f ᵀ

In deep learning, the Jacobian of a layer's output w.r.t. its input describes how that layer transforms gradients during backprop. For a linear layer z = Wx + b, the Jacobian w.r.t. x is simply W.

DIAGRAM

Multivariable chain rule example:
  f(g,h) = g²h + h,  g(x)=3x,  h(x)=x²

  ∂f/∂g = 2gh
  ∂f/∂h = g² + 1
  dg/dx = 3
  dh/dx = 2x

  df/dx = (∂f/∂g)(dg/dx) + (∂f/∂h)(dh/dx)
        = 2gh·3 + (g²+1)·2x
        = 6gh + 2x(g²+1)

  Jacobian of f: R² → R²
  f₁(x,y) = x² + y
  f₂(x,y) = xy
        ∂f₁/∂x  ∂f₁/∂y     2x  1
  J = [               ] = [      ]
        ∂f₂/∂x  ∂f₂/∂y      y  x

CODE

PYTHON
1import numpy as np
2from sympy import *
3import torch
4
5x, y, g, h = symbols('x y g h')
6
7# ── Multivariable chain rule: Raschka's example ───────────────────
8print("Multivariable chain rule:")
9print("f(g,h) = g²h + h,  g(x)=3x,  h(x)=x²")
10
11f_gh = g**2 * h + h
12g_x  = 3*x
13h_x  = x**2
14
15df_dg  = diff(f_gh, g)   # 2gh
16df_dh  = diff(f_gh, h)   # g^2 + 1
17dg_dx  = diff(g_x, x)    # 3
18dh_dx  = diff(h_x, x)    # 2x
19
20# Substitute g(x), h(x) back
21df_dg_x = df_dg.subs([(g, g_x), (h, h_x)])
22df_dh_x = df_dh.subs([(g, g_x), (h, h_x)])
23
24df_dx = df_dg_x * dg_dx + df_dh_x * dh_dx
25print(f"  df/dx = {simplify(df_dx)}")
26
27# ── Jacobian matrix ────────────────────────────────────────────────
28print("
29Jacobian of f: R² → R²")
30x_sym, y_sym = symbols('x y')
31f1 = x_sym**2 + y_sym
32f2 = x_sym * y_sym
33
34J = Matrix([[diff(f1, x_sym), diff(f1, y_sym)],
35            [diff(f2, x_sym), diff(f2, y_sym)]])
36print(f"  J = {J}")
37print(f"  At (2,3): J = {J.subs([(x_sym,2),(y_sym,3)])}")
38
39# ── Jacobian with PyTorch autograd ────────────────────────────────
40print("
41Jacobian via PyTorch:")
42def f_torch(x):
43    return torch.stack([x[0]**2 + x[1], x[0] * x[1]])
44
45x_pt = torch.tensor([2.0, 3.0], requires_grad=True)
46J_pt = torch.autograd.functional.jacobian(f_torch, x_pt)
47print(f"  J = {J_pt}")
48
49# ── Chain rule in neural network layer ───────────────────────────
50print("
51Jacobian of linear layer z = Wx:")
52W = np.array([[1, 2, 3],
53              [4, 5, 6]])    # (2, 3) — transforms R³ → R²
54x_input = np.array([1.0, 2.0, 3.0])
55z = W @ x_input
56
57# Jacobian of z w.r.t. x is simply W!
58print(f"  Input shape: {x_input.shape}")
59print(f"  Output shape: {z.shape}")
60print(f"  Jacobian dz/dx = W:
61{W}")
62print("  During backprop: dL/dx = Wᵀ · dL/dz")
← PREV6. Second Order & HessianNEXT →8. Integration