| """The worked numbers of Chapter 6, asserted against the Lecture 6 files. |
| |
| Every number checked here appears in the chapter text, and the file under test computes each one |
| and stores none of them. We never run the training loops, and the tests call the closed-form and |
| numerical functions the lectures build their reports from. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import importlib.util |
| import math |
| import sys |
| from pathlib import Path |
|
|
| import pytest |
| import torch |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(ROOT)) |
|
|
|
|
| def load(name: str): |
| """Import one lecture file by path, so that ``lecture_6`` needs no package marker.""" |
| spec = importlib.util.spec_from_file_location(f"lecture_6_{name}", |
| ROOT / "lecture_6" / f"{name}.py") |
| module = importlib.util.module_from_spec(spec) |
| spec.loader.exec_module(module) |
| return module |
|
|
|
|
| def close(values, expected, tolerance=5e-4): |
| """Compare a list of numbers against the chapter's, coordinate by coordinate.""" |
| assert len(values) == len(expected) |
| for got, want in zip(values, expected): |
| assert got == pytest.approx(want, abs=tolerance) |
|
|
|
|
| |
| def test_flow_map_listing(): |
| flow_map = load("flow_map") |
| x = torch.tensor(0.7, dtype=torch.float64) |
| assert float(flow_map.flow(x, 0.0, 1.0)) == pytest.approx(4.1, abs=1e-9) |
| assert float(flow_map.flow(flow_map.flow(x, 0.0, 0.5), 0.5, 1.0)) == pytest.approx(4.1, abs=1e-9) |
| assert float(flow_map.flow(flow_map.flow(x, 0.0, 1.0), 1.0, 0.0)) == pytest.approx(0.7, abs=1e-9) |
| from torch.func import jacrev |
|
|
| assert float(jacrev(lambda z: flow_map.flow(z, 0.0, 1.0))(x)) == pytest.approx(3.0, abs=1e-9) |
|
|
|
|
| def test_flow_map_axioms_and_liouville(): |
| flow_map = load("flow_map") |
| x = torch.tensor(0.7, dtype=torch.float64) |
| identity, cocycle, tangent = flow_map.check_axioms(x, 200) |
| assert identity < 1e-12 and cocycle < 1e-10 and tangent < 1e-4 |
| assert flow_map.liouville_log_det(0.0, 1.0, 4000) == pytest.approx(math.log(3.0), abs=1e-6) |
| assert float(flow_map.jacobian_by_ode(x, 0.0, 1.0, 4000)) == pytest.approx(3.0, abs=1e-3) |
|
|
|
|
| def test_flow_map_transport_and_likelihood(): |
| flow_map = load("flow_map") |
| source = flow_map.normal_density(0.7, 0.0, 1.0) |
| data = flow_map.normal_density(4.1, 2.0, 3.0) |
| assert source == pytest.approx(0.3123, abs=1e-4) |
| assert data == pytest.approx(0.1041, abs=1e-4) |
| assert source / data == pytest.approx(3.0, abs=1e-6) |
| assert math.log(source) - math.log(3.0) == pytest.approx(-2.2625, abs=1e-4) |
| assert math.log(data) == pytest.approx(-2.2625, abs=1e-4) |
|
|
|
|
| |
| def test_stay_still_residuals(): |
| residuals = load("residuals") |
| report = residuals.stay_still_residuals() |
| assert report["stay-still Lagrangian"] == pytest.approx(1.0, abs=1e-9) |
| assert report["stay-still Eulerian"] == pytest.approx(1.0, abs=1e-9) |
| assert report["stay-still composition"] == pytest.approx(0.0, abs=1e-12) |
| assert report["stay-still diagonal gap"] == pytest.approx(0.0, abs=1e-12) |
|
|
|
|
| def test_identity_axiom_is_free(): |
| residuals = load("residuals") |
| torch.manual_seed(0) |
| model = residuals.Map(dim=2, width=16) |
| x, t = torch.randn(8, 2), torch.rand(8) |
| assert torch.allclose(model.X(x, t, t), x) |
|
|
|
|
| |
| def test_meanflow_scalar_identity(): |
| meanflow = load("meanflow") |
| report = meanflow.scalar_check() |
| assert report["v_0,1(1)"] == pytest.approx(math.e - 1.0, abs=1e-5) |
| assert report["b(1, 0)"] == pytest.approx(1.0, abs=1e-6) |
| assert report["partial_s v"] == pytest.approx(-1.0, abs=1e-5) |
| assert report["grad_x v_0,1"] == pytest.approx(math.e - 1.0, abs=1e-5) |
| assert report["d/ds v_0,1"] == pytest.approx(math.e - 2.0, abs=1e-5) |
| assert report["identity gap"] == pytest.approx(0.0, abs=1e-5) |
| assert report["one map evaluation"] == pytest.approx(math.e, abs=1e-5) |
| assert report["one Euler step"] == pytest.approx(2.0, abs=1e-6) |
|
|
|
|
| def test_meanflow_correction_vanishes_on_the_gaussian_flow(): |
| meanflow = load("meanflow") |
| report = meanflow.gaussian_check() |
| assert report["v_0,t(0.7) spread"] == pytest.approx(0.0, abs=1e-6) |
| assert report["v_0,1(0.7)"] == pytest.approx(3.4, abs=1e-5) |
| assert report["b(0.7, 0)"] == pytest.approx(3.4, abs=1e-5) |
|
|
|
|
| |
| def test_shortcut_doubling(): |
| shortcut = load("shortcut") |
| report = shortcut.scalar_doubling() |
| assert report["s(1, 0, d)"] == pytest.approx(1.136102, abs=1e-6) |
| assert report["state after the first step"] == pytest.approx(1.284025, abs=1e-6) |
| assert report["s(x', d, d)"] == pytest.approx(1.458783, abs=1e-6) |
| assert report["average of the two"] == pytest.approx(1.297443, abs=1e-6) |
| assert report["s(1, 0, 2d) direct"] == pytest.approx(1.297443, abs=1e-6) |
| assert report["doubling gap"] == pytest.approx(0.0, abs=1e-12) |
| assert report["state after the doubled step"] == pytest.approx(math.exp(0.5), abs=1e-9) |
| assert shortcut.ladder()[0] == 1.0 / 128.0 and shortcut.ladder()[-1] == 1.0 |
|
|
|
|
| |
| def test_consistency_boundary_parameterization(): |
| consistency = load("consistency") |
| report = consistency.scalar_consistency() |
| assert report["f(0.5)"] == pytest.approx(math.exp(0.5), abs=1e-6) |
| assert report["f(0.9)"] == pytest.approx(math.exp(0.1), abs=1e-6) |
| assert report["f(1.0)"] == pytest.approx(1.0, abs=1e-9) |
| assert report["c_out(1.0)"] == pytest.approx(0.0, abs=1e-12) |
| assert report["v_0.5,1(1)"] == pytest.approx(1.297443, abs=1e-6) |
|
|
|
|
| def test_consistency_teacher_floor(): |
| consistency = load("consistency") |
| report = consistency.teacher_floor() |
| assert report["f(1, 0.5)"] == pytest.approx(1.648721, abs=1e-6) |
| assert report["f(1.01, 0.51)"] == pytest.approx(1.648639, abs=1e-6) |
| assert report["distillation residual"] == pytest.approx(8.189e-5, rel=1e-3) |
| assert report["local Euler error"] == pytest.approx(5.0167e-5, rel=1e-3) |
| assert report["carried to the data end"] == pytest.approx( |
| report["distillation residual"], rel=1e-6) |
|
|
|
|
| def test_consistency_trajectory_head(): |
| consistency = load("consistency") |
| report = consistency.scalar_ctm() |
| assert report["a_0.0,0.5"] == pytest.approx(0.5, abs=1e-12) |
| assert report["g_0.0,0.5"] == pytest.approx(2.297443, abs=1e-6) |
| assert report["a_0.0,1.0"] == pytest.approx(0.0, abs=1e-12) |
| assert report["g_0.0,1.0"] == pytest.approx(math.e, abs=1e-6) |
| assert report["two routes apart"] == pytest.approx(0.0, abs=1e-12) |
|
|
|
|
| |
| def test_fmlm_denoiser_cocycle(): |
| fmlm = load("fmlm") |
| report = fmlm.scalar_cocycle() |
| assert report["gamma"] == pytest.approx(0.2, abs=1e-12) |
| assert report["delta_s,t"] == pytest.approx(2.096158, abs=1e-6) |
| assert report["delta_s,u"] == pytest.approx(1.932957, abs=1e-6) |
| assert report["delta_u,t"] == pytest.approx(2.136959, abs=1e-6) |
| assert report["two legs combined"] == pytest.approx(2.096158, abs=1e-6) |
| assert report["convex weights"] == (0.25, 0.75) |
| assert report["convex form"] == pytest.approx(math.exp(0.6), abs=1e-6) |
| assert report["X_s,t(x)"] == pytest.approx(math.exp(0.6), abs=1e-9) |
|
|
|
|
| def test_fmlm_worked_kl(): |
| fmlm = load("fmlm") |
| report = fmlm.worked_kl() |
| assert report["worked gamma"] == pytest.approx(0.2, abs=1e-12) |
| close(report["composed target"], [0.06, 0.12, 0.76, 0.06], 1e-6) |
| assert report["target mass"] == pytest.approx(1.0, abs=1e-6) |
| close(report["the four terms"], [-0.030650, -0.026777, 0.118823, -0.030650], 1e-6) |
| assert report["worked KL"] == pytest.approx(0.030747, abs=1e-6) |
|
|
|
|
| def test_fmlm_gradient_comparison(): |
| """Section 6.27.4, the two gradients in the logits at the chapter's target and student.""" |
| fmlm = load("fmlm") |
| report = fmlm.gradient_comparison() |
| close(report["KL gradient"], [0.04, 0.03, -0.11, 0.04], 1e-6) |
| close(report["squared-error gradient"], [0.0099, 0.01335, -0.03315, 0.0099], 1e-6) |
| close(report["squared over KL"], [0.2475, 0.445, 0.3014, 0.2475], 1e-4) |
| close(report["saturated student"], [1e-6, 0.1666665, 0.7222215, 0.111111], 1e-6) |
| assert report["saturated KL gradient"] == pytest.approx(-0.0600, abs=1e-4) |
| |
| |
| |
| |
| assert report["saturated squared-error gradient, in 1e-08"] == pytest.approx(-4.6171, abs=1e-3) |
| assert report["orders of magnitude between them"] == pytest.approx(6.11, abs=0.01) |
|
|
|
|
| def test_fmlm_decoding_clock(): |
| fmlm = load("fmlm") |
| tau, start = fmlm.clock(fmlm.BOOK_VOCAB) |
| assert start == pytest.approx(0.999967, abs=1e-5) |
| assert tau(0.60) == pytest.approx(0.006, abs=5e-4) |
| assert tau(0.75) == pytest.approx(0.141, abs=5e-4) |
| assert tau(0.80) == pytest.approx(0.456, abs=1e-3) |
| assert tau(0.85) == pytest.approx(0.931, abs=1e-3) |
| grid = fmlm.even_in_tau(fmlm.BOOK_VOCAB) |
| close(grid, [0.0, 0.773490, 0.804481, 0.828095, 1.0], 5e-5) |
| close([tau(t) for t in (0.25, 0.5, 0.75)], [0.00009, 0.00124, 0.14109], 1e-5) |
|
|
|
|
| def test_fmlm_clock_weights(): |
| fmlm = load("fmlm") |
| rounded = [0.0, 0.773, 0.804, 0.828, 1.0] |
| pairs = [((1 - b) / (1 - a), (b - a) / (1 - a)) for a, b in zip(rounded[:-1], rounded[1:])] |
| close([p[0] for p in pairs], [0.2270, 0.8634, 0.8776, 0.0], 1e-4) |
| close([p[1] for p in pairs], [0.7730, 0.1366, 0.1224, 1.0], 1e-4) |
| times = [torch.tensor([v]) for v in rounded[1:4]] |
| assert float(fmlm.composition_weight(*times)) == pytest.approx(0.494620, abs=1e-5) |
| early = [torch.tensor([v]) for v in (0.0, 0.773, 0.804)] |
| assert float(fmlm.composition_weight(*early)) == pytest.approx(0.830144, abs=1e-5) |
|
|
|
|
| def test_fmlm_endpoint_slice_is_trained(): |
| """Section 6.26 asks for t = 1 on purpose, since a continuous draw never lands there.""" |
| fmlm = load("fmlm") |
| assert 0.0 < fmlm.ENDPOINT_FRACTION < 1.0 |
| tokens, _ = fmlm.synthetic_dna(num_sequences=64, length=16, seed=0) |
| model = fmlm.TwoTimeMLP(dim=16 * 4, width=16) |
| torch.manual_seed(0) |
| loss_fn = fmlm.make_loss(model, tokens, 16, 4, 256, torch.device("cpu")) |
| assert torch.isfinite(loss_fn(0)) |
| |
| x = torch.randn(2, 16, 4) |
| s, t = torch.zeros(2), torch.ones(2) |
| prediction = fmlm.predict(model, x, s, t, 16, 4) |
| assert torch.allclose(fmlm.step_map(x, prediction, s, t), prediction) |
|
|
|
|
| |
| def test_categorical_worked_cancellation(): |
| categorical = load("categorical_map") |
| report = categorical.worked_cancellation() |
| close(report["state x"], [-0.2591, -0.6675, -0.0467], 1e-4) |
| close(report["X_s,t(x)"], [0.1611, -0.2420, 0.0943], 1e-4) |
| close(report["pi_s,t(x)"], [0.5813, 0.1835, 0.2352], 1e-4) |
| close(report["pi_t,t(X_s,t(x))"], [0.6927, 0.0917, 0.2157], 1e-4) |
| assert report["gamma_s,t"] == pytest.approx(0.2, abs=1e-12) |
| close(report["d pi / dt"], [0.5570, -0.4592, -0.0979], 1e-4) |
| close(report["central difference"], [0.5570, -0.4592, -0.0979], 1e-4) |
| close(report["gap a"], [-0.111, 0.092, 0.020], 1e-3) |
| close(report["drift b"], [0.111, -0.092, -0.020], 1e-3) |
|
|
|
|
| def test_categorical_bound_is_positive_where_the_residual_is_zero(): |
| categorical = load("categorical_map") |
| report = categorical.worked_cancellation() |
| assert report["residual |a + b|^2"] < 1e-20 |
| assert report["4 KL"] == pytest.approx(0.156, abs=1e-3) |
| assert report["2 |b|^2"] == pytest.approx(0.042, abs=1e-3) |
| assert report["the bound"] == pytest.approx(0.199, abs=2e-3) |
| assert report["teacher entropy"] == pytest.approx(0.8043, abs=1e-3) |
| assert report["cross-entropy"] == pytest.approx(0.8434, abs=1e-3) |
| assert report["cross-entropy"] == pytest.approx( |
| 0.25 * report["4 KL"] + report["teacher entropy"], abs=1e-6) |
|
|
|
|
| def test_categorical_averaged_bound(): |
| categorical = load("categorical_map") |
| report = categorical.averaged_bound(1000, seed=0) |
| assert report["averaged residual"] < 1e-18 |
| assert report["averaged 4 KL"] == pytest.approx(0.39, abs=0.06) |
| assert report["averaged 2 |b|^2"] == pytest.approx(0.095, abs=0.03) |
|
|
|
|
| |
| def test_discrete_derivative_target_leaves_the_simplex(): |
| discrete = load("discrete_map") |
| report = discrete.worked_derivative() |
| assert report["derivative prefactor"] == pytest.approx(6.6667, abs=1e-3) |
| close(report["derivative target"], [-0.3333, -0.6667, 1.3333, -0.3333], 1e-4) |
| assert report["target sum"] == pytest.approx(0.0, abs=1e-6) |
| assert report["negative entries"] == 3 |
|
|
|
|
| def test_discrete_logit_repair(): |
| discrete = load("discrete_map") |
| report = discrete.worked_repair() |
| assert report["weighted average of dz/dt"] == pytest.approx(0.0, abs=1e-6) |
| close(report["d psi/dt"], [1.92, -1.14, -0.78], 1e-4) |
| close(report["additive target"], [-0.284, 0.828, 0.456], 1e-4) |
| assert report["additive target sum"] == pytest.approx(1.0, abs=1e-6) |
| assert report["additive negative entries"] == 1 |
| close(report["multiplicative factor"], [2.920, 0.544, 0.480], 1e-4) |
| close(report["quotient"], [0.0342, 1.1029, 0.6250], 1e-4) |
| assert report["quotient mass"] == pytest.approx(1.7622, abs=1e-4) |
| close(report["repaired target"], [0.0194, 0.6259, 0.3547], 1e-4) |
| close(report["log of the factor"], [1.0716, -0.6088, -0.7340], 1e-4) |
| assert report["gamma_s,t"] == pytest.approx(0.2, abs=1e-12) |
| assert report["gamma^E_s,t"] == pytest.approx(0.8, abs=1e-12) |
|
|
|
|
| |
| def test_stochastic_weak_floor(): |
| stochastic = load("stochastic_map") |
| report = stochastic.weak_floor(draws=200_000, seed=0) |
| assert report["K_1"] == pytest.approx(3.194528, abs=1e-5) |
| assert report["sqrt K_1"] == pytest.approx(1.7873, abs=1e-4) |
| assert report["sqrt K_0.5"] == pytest.approx(0.9269, abs=1e-4) |
| assert report["one full step"] == pytest.approx(2.2453, abs=1e-3) |
| assert report["two half steps"] == pytest.approx(0.4673, abs=1e-3) |
| assert report["gap"] == pytest.approx(1.778, abs=1e-3) |
| assert report["squared gap"] == pytest.approx(3.161, abs=1e-3) |
| assert report["floor 2 K_1"] == pytest.approx(math.e ** 2 - 1.0, abs=1e-9) |
| assert report["carried variance equals K_1"] == pytest.approx(report["K_1"], abs=1e-9) |
| assert report["the draw that would close it"] == pytest.approx(-0.195, abs=1e-3) |
| assert report["Monte Carlo floor"] == pytest.approx(6.389, abs=0.05) |
|
|
|
|
| def test_stochastic_chen_relations(): |
| stochastic = load("stochastic_map") |
| report = stochastic.worked_chen() |
| assert report["I0 over [0,1]"] == pytest.approx(0.50, abs=1e-9) |
| assert report["I1 over [0,1]"] == pytest.approx(-0.20, abs=1e-9) |
| assert report["fused variance at h=1"] == pytest.approx(1.0 / 3.0, abs=1e-9) |
|
|
|
|
| def test_stochastic_truncation(): |
| stochastic = load("stochastic_map") |
| report = stochastic.truncation_error(paths=512, grid=2048, seed=0) |
| for n, expected in enumerate((1.0, 1.0 / 3.0, 1.0 / 5.0)): |
| assert report[f"coefficient variance I{n}"] == pytest.approx(expected, rel=0.15) |
| errors = [report[f"pathwise error with {k}"] for k in (1, 2, 3)] |
| assert errors[0] > 5.0 * errors[1] > 25.0 * errors[2] |
|
|
|
|
| |
| def test_expanding_local_clocks(): |
| expanding = load("expanding") |
| report = expanding.worked_clocks() |
| close(report["local clocks at t=0.8"], [0.800, 0.800, 0.692, 0.333], 1e-3) |
| close(report["clock speeds"], [1.0, 1.0, 1.538, 3.333], 1e-3) |
| close(report["source share"], [0.200, 0.200, 0.308, 0.667], 1e-3) |
| assert report["last clock from its speed"] == pytest.approx(0.667, abs=1e-3) |
| assert report["live coordinates at 0.8"] == 4 and report["live coordinates at 0.6"] == 3 |
|
|
|
|
| def test_expanding_composition_weight(): |
| expanding = load("expanding") |
| report = expanding.worked_omega() |
| assert report["omega at (0.3, 0.6, 0.9)"] == pytest.approx(0.125, abs=1e-6) |
| assert report["omega for a birth at 0.0"] == pytest.approx(0.125, abs=1e-6) |
| assert report["omega for a birth at 0.7"] == pytest.approx(0.0, abs=1e-9) |
| assert 0.0 < report["omega for a birth at 0.35"] < report["omega for a birth at 0.0"] |
|
|
|
|
| |
| def test_posterior_tilt_and_jensen_gap(): |
| posterior = load("posterior_map") |
| report = posterior.two_point() |
| close(report["tilted target"], [0.1192, 0.8808], 1e-4) |
| close(report["posterior at x=0.7"], [0.9820, 0.0180], 1e-4) |
| close(report["tilted posterior"], [0.8808, 0.1192], 1e-4) |
| assert report["V_t(0.7)"] == pytest.approx(0.1088, abs=1e-4) |
| assert report["V over the data"] == pytest.approx(1.4338, abs=1e-4) |
| assert report["single-endpoint estimate"] == pytest.approx(1.0, abs=1e-9) |
| assert report["Jensen gap"] == pytest.approx(0.434, abs=1e-3) |
|
|
|
|
| def test_posterior_value_bias_and_gradient(): |
| posterior = load("posterior_map") |
| report = posterior.value_bias() |
| close([report[f"E[V_{k}]"] for k in (1, 2, 4, 8)], [1.000, 1.217, 1.340, 1.393], 1e-3) |
| close([report[f"bias at K={k}"] for k in (1, 2, 4, 8)], [0.434, 0.217, 0.093, 0.041], 1e-3) |
| weights, weighted, plain = posterior.value_gradient( |
| torch.tensor([0.2, 1.5, 0.9]), torch.tensor([0.4, -1.2, 0.5])) |
| close(weights.tolist(), [0.1496, 0.5491, 0.3013], 1e-4) |
| assert weighted == pytest.approx(-0.4483, abs=1e-4) |
| assert plain == pytest.approx(-0.1000, abs=1e-6) |
|
|
|
|
| def test_posterior_renoising_and_meta_fixed_point(): |
| posterior = load("posterior_map") |
| report = posterior.renoise_example() |
| assert report["renoise scale"] == pytest.approx(0.75, abs=1e-9) |
| assert report["renoise variance"] == pytest.approx(0.1375, abs=1e-9) |
| assert report["renoise standard deviation"] == pytest.approx(0.3708, abs=1e-4) |
| assert report["proposal weight on the far mode"] == pytest.approx(0.420, abs=1e-3) |
| assert report["penalty at z=2"] == pytest.approx(10.125, abs=1e-6) |
| assert report["penalty at z=0"] == pytest.approx(6.125, abs=1e-6) |
| assert report["corrected weight on the far mode"] == pytest.approx(0.0131, abs=1e-4) |
| meta = posterior.stop_gradient_fixed_point(torch.tensor([1.0, 2.0, 5.0]), |
| torch.tensor([0.4, 1.0, 3.6])) |
| assert meta["E[w]"] == pytest.approx(8.0 / 3.0, abs=1e-6) |
| assert meta["E[G]"] == pytest.approx(5.0 / 3.0, abs=1e-6) |
| assert meta["stop-gradient fixed point"] == pytest.approx(0.625, abs=1e-6) |
| assert meta["per-sample ratio average"] == pytest.approx(0.540, abs=1e-6) |
| assert meta["gap"] == pytest.approx(-0.085, abs=1e-3) |
| assert meta["covariance of G/w and w"] == pytest.approx(0.2267, abs=1e-4) |
| assert meta["gap from the covariance"] == pytest.approx(meta["gap"], abs=1e-6) |
|
|
