Lab U3: Jacobians and local linearization

Unit: Unit 3, Beyond Linear Maps
Role: Required
Textbook sections: Jacobians; local linearization; chain rule as matrix multiplication
Core path: symbolic Jacobians, evaluating Jacobians, actual versus local linear prediction, chain-rule shape checks
This lab practices reading Jacobian computations. The main comparison is

[ F(\mathbf{a}+\mathbf{h}) \quad \text{versus} \quad F(\mathbf{a}) + J_F(\mathbf{a})\mathbf{h}. ]

The first quantity is the actual nonlinear output. The second is the local linear prediction.

Computational tools used in this lab

Before starting, review these parts of Appendix B, NumPy and SymPy Quick Reference for the Labs:

• Appendix B.2: NumPy arrays, vectors, matrices, and shapes
• Appendix B.4: elementwise arithmetic versus linear algebra
• Appendix B.9: numerical checks
• Appendix B.10: NumPy linear algebra commands
• Appendix B.11: SymPy symbolic computation

The goal is to interpret the mathematical computation, not to memorize every command.

Part 0. Setup and shape habits

Shape habit: Before reading a product such as J_at_a @ h, identify the shape of J_at_a and the length of h.

Predict before running:

1. What is the input dimension?
2. What is the output dimension?
3. Why is J_example @ h_example defined?
4. What is the shape of the output?

import numpy as np
import sympy as sp

np.set_printoptions(precision=4, suppress=True)

J_example = np.array([[1.0, 2.0],
                      [0.0, -1.0],
                      [3.0, 4.0]])
h_example = np.array([0.1, -0.2])

J_example.shape, h_example.shape, J_example @ h_example

Part 1. Core path: symbolic Jacobian

Math reminder:

[ F(x,y)= \begin{bmatrix} x^2y\\ xe^y \end{bmatrix}. ]

Predict before running:

1. How many input variables does (F) have?
2. How many output components does (F) have?
3. What should the shape of (J) be?

x, y = sp.symbols("x y")

F = sp.Matrix([
    x**2 * y,
    x * sp.exp(y)
])

J = F.jacobian([x, y])
F, J

Interpretation check:

1. Which row corresponds to the second component of (F)?
2. Which column describes changing (y)?
3. How does this match the Unit 1 idea "rows measure; columns contribute"?

Part 2. Evaluate a Jacobian at a point

Math reminder:

[ \mathbf{a}=\begin{bmatrix}1\\0\end{bmatrix}. ]

Predict before running:

1. Which symbols should be replaced by the entries of (\mathbf{a})?
2. Should (J_F(\mathbf{a})) still have the same shape as (J_F(x,y))?

a = np.array([1.0, 0.0])

J_at_a_sym = J.subs({x: a[0], y: a[1]})
J_at_a = np.array(J_at_a_sym, dtype=float)

J_at_a_sym, J_at_a, J_at_a.shape

Interpretation check:

1. What is (J_F(\mathbf{a}))?
2. What is its shape?
3. What does row 1 measure?
4. What does column 2 measure?

Part 3. Actual output versus local linear prediction

Math reminder: compare the actual nonlinear output with the local linear prediction.

[ F(\mathbf{a}+\mathbf{h}) \quad \text{versus} \quad F(\mathbf{a})+J_F(\mathbf{a})\mathbf{h}. ]

Predict before running:

1. Which line below will compute the actual nonlinear output?
2. Which line below will compute the local linear prediction?

def F_np(v):
    x_val, y_val = v
    return np.array([
        x_val**2 * y_val,
        x_val * np.exp(y_val)
    ], dtype=float)

h = np.array([0.1, -0.2])

actual = F_np(a + h)
linear = F_np(a) + J_at_a @ h
error = actual - linear

actual, linear, error

Interpretation check:

1. Which output is the actual nonlinear output?
2. Which output is the local linear prediction?
3. What does the error vector measure?
4. Is the prediction exact?

Part 4. What happens when the input change gets smaller?

Predict before running: as the scale decreases, what should happen to the error norm?

for scale in [1.0, 0.5, 0.25, 0.125]:
    h_scaled = scale * h
    actual_scaled = F_np(a + h_scaled)
    linear_scaled = F_np(a) + J_at_a @ h_scaled
    error_scaled = actual_scaled - linear_scaled
    print(scale, actual_scaled, linear_scaled, error_scaled, np.linalg.norm(error_scaled))

Interpretation check:

1. As the scale decreases, what happens to the error norm?
2. Why does this support the phrase "local linear approximation"?
3. Why would this not prove that the approximation is good far away from (\mathbf{a})?

Part 5. Scalar local linearization

Math reminder: for a scalar-valued function, the gradient gives the local linear part.

[ f(x,y)=x^2+xy. ]

Predict before running:

1. Is the output of (f) a scalar or a vector?
2. What shape should the gradient have at a point?

f = x**2 + x*y
grad_f = sp.Matrix([sp.diff(f, x), sp.diff(f, y)])

a_scalar = np.array([1.0, 2.0])
h_scalar = np.array([0.1, -0.2])

grad_at_a = np.array(
    grad_f.subs({x: a_scalar[0], y: a_scalar[1]}),
    dtype=float
).reshape(2)

def f_np(v):
    x_val, y_val = v
    return x_val**2 + x_val*y_val

prediction = f_np(a_scalar) + grad_at_a @ h_scalar
actual_scalar = f_np(a_scalar + h_scalar)

grad_f, grad_at_a, prediction, actual_scalar, actual_scalar - prediction

Interpretation check:

1. Why is the gradient a vector but the output of (f) is a scalar?
2. Where does the dot product appear?
3. What does the error measure?
4. How does this connect to tangent planes?

Part 6. Affine maps are exact local models

Math reminder: an affine map has the form A @ v + b. It is not usually linear as a function of v, but its change is controlled by the linear map A.

Predict before running: what should the error be?

A = np.array([[2.0, -1.0],
              [0.0,  3.0]])
b = np.array([1.0, -2.0])

def affine_map(v):
    return A @ v + b

a_lin = np.array([1.0, 2.0])
h_lin = np.array([0.3, -0.4])

actual_affine = affine_map(a_lin + h_lin)
linear_prediction_affine = affine_map(a_lin) + A @ h_lin

actual_affine, linear_prediction_affine, actual_affine - linear_prediction_affine

Interpretation check:

1. Why is the error zero?
2. Is v -> A @ v + b linear or affine?
3. What is the linear map on input changes?
4. Why does the bias vector disappear when comparing changes?

Part 7. Jacobian columns

Math reminder: a column of a Jacobian describes the predicted output change from changing one input coordinate.

Predict before running: what should J_at_a @ e1 and J_at_a @ e2 return?

e1 = np.array([1.0, 0.0])
e2 = np.array([0.0, 1.0])

J_at_a @ e1, J_at_a @ e2, J_at_a[:, 0], J_at_a[:, 1]

Interpretation check:

1. What does J_at_a @ e1 return?
2. What does J_at_a @ e2 return?
3. Why do these match the columns of J_at_a?
4. What does each column measure locally?

Part 8. Locally forgotten directions

Extension connection: Unit 2 null-space language helps interpret a Jacobian that sends a nonzero input change to zero.

Predict before running: what should happen to J_forget @ h_forget?

J_forget = np.array([[1.0, 2.0],
                     [2.0, 4.0]])

h_forget = np.array([-2.0, 1.0])

J_forget @ h_forget

Interpretation check:

1. Why is h_forget a locally forgotten direction?
2. How does this connect to Unit 2 null spaces?
3. Does (J_F(\mathbf{a})\mathbf{h}=\mathbf{0}) prove that (F(\mathbf{a}+\mathbf{h})=F(\mathbf{a})) exactly?
4. What phrase should be used instead: exact equality or first-order prediction?

Part 9. Chain rule as matrix multiplication

Math reminder:

[ G:\mathbb R^2\to\mathbb R^3, \qquad F:\mathbb R^3\to\mathbb R^2. ]

Before multiplying Jacobians, check the input and output dimensions of each map.

u, v = sp.symbols("u v")
s, t, r = sp.symbols("s t r")

G = sp.Matrix([
    u + v,
    u - v,
    u*v
])

F_outer = sp.Matrix([
    s**2 + t,
    sp.exp(r) + s*t
])

JG = G.jacobian([u, v])
JF_outer = F_outer.jacobian([s, t, r])

JG, JF_outer

base = {u: 1, v: 2}

G_at_base = G.subs(base)

JF_at_G = JF_outer.subs({
    s: G_at_base[0],
    t: G_at_base[1],
    r: G_at_base[2]
})

JG_at_base = JG.subs(base)

J_comp = JF_at_G @ JG_at_base

G_at_base, JF_at_G, JG_at_base, J_comp

Interpretation check:

1. What is the shape of JG?
2. What is the shape of JF_at_G?
3. Which Jacobian acts first on an input change?
4. Why is the product JF_at_G @ JG_at_base?
5. What shape should the composite Jacobian have?

# Debug check: uncomment to test the wrong order.
# JG_at_base @ JF_at_G

Debug check: Why should the reversed product fail or give the wrong interpretation?

Part 10. Finite-difference check for the chain rule

Predict before running: for a small input change, should the chain-rule linear prediction be close to the actual composite output?

def G_np(v_in):
    u_val, v_val = v_in
    return np.array([
        u_val + v_val,
        u_val - v_val,
        u_val * v_val
    ], dtype=float)

def F_outer_np(w_in):
    s_val, t_val, r_val = w_in
    return np.array([
        s_val**2 + t_val,
        np.exp(r_val) + s_val*t_val
    ], dtype=float)

def composite_np(v_in):
    return F_outer_np(G_np(v_in))

base_np = np.array([1.0, 2.0])
h_chain = np.array([0.01, -0.02])

J_comp_np = np.array(J_comp, dtype=float)

actual_chain = composite_np(base_np + h_chain)
linear_chain = composite_np(base_np) + J_comp_np @ h_chain

actual_chain, linear_chain, actual_chain - linear_chain

Interpretation check:

1. Which line computes the actual composite output?
2. Which line computes the local linear prediction?
3. Where does the chain rule matrix product appear?
4. Why is h_chain chosen small?

Part 11. Reading a tiny nonlinear block

Use hidden instead of h in code so it does not get confused with the input-change vector (\mathbf{h}).

Predict before running: which lines are affine, and which line is nonlinear?

def relu(z):
    return np.maximum(z, 0)

W1 = np.array([[1.0, -1.0],
               [0.5,  2.0],
               [-1.0, 1.0]])

b1 = np.array([0.1, -0.2, 0.3])

W2 = np.array([[2.0, 0.0, -1.0]])

b2 = np.array([0.5])

x_input = np.array([1.0, 2.0])

z1 = W1 @ x_input + b1
hidden = relu(z1)
z2 = W2 @ hidden + b2

z1, hidden, z2

Interpretation check:

1. Which lines are affine?
2. Which line is nonlinear?
3. Which variables are vectors?
4. Which objects are matrices?
5. Why is the whole block usually not linear?
6. Where would Jacobians enter if we wanted a local linear approximation?

Part 12. Review code-reading bank

These are short checks for discussion or self-review. They are meant to be answered without writing long programs.

Exam-style check

1. Shape check. In J.shape, which entry gives the output dimension and which gives the input dimension?
2. Predicted change. In J_at_a @ h, is the result an actual nonlinear output or a predicted output change?
3. Actual output. In F_np(a + h), is this the actual nonlinear output or the local prediction?
4. Affine versus linear. In F_np(a) + J_at_a @ h, why is the expression affine in a + h but linear in h?
5. Error vector. What does actual - linear measure?
6. Column reading. In J_at_a[:, 0], which coordinate direction does this column describe?
7. Scalar local model. In grad_at_a @ h_scalar, why does a dot product appear for scalar-valued functions?
8. Chain-rule order. In JF_at_G @ JG_at_base, which map is applied first?
9. Wrong order. What is wrong with trying JG_at_base @ JF_at_G?
10. Affine exactness. Why is affine_map(a_lin+h_lin) - (affine_map(a_lin)+A @ h_lin) the zero vector?
11. Forgotten direction. What does J_forget @ h_forget being zero mean locally?
12. Tiny block. Is W1 @ x_input + b1 linear or affine as a function of x_input?

Answer sketches

1. For a Jacobian with shape (m, n), m is the output dimension and n is the input dimension.
2. A predicted output change.
3. The actual nonlinear output.
4. The base value is a shift; the input-change part J_at_a @ h is linear in h.
5. The difference between the actual nonlinear output and the local linear prediction.
6. The first input-coordinate direction.
7. A scalar-valued local model sends an input change to one number, so the gradient dots with the input change.
8. JG_at_base acts first on the input change; then JF_at_G acts on the resulting change.
9. The dimensions do not match the order of composition, and the interpretation is reversed.
10. For an affine map, changes are exactly controlled by the linear part A.
11. The first-order prediction says that direction produces no output change.
12. Affine, because the bias vector is added after multiplication.