path: root/experiments/conv_local_smoke.py

#!/usr/bin/env python3
"""Prove the convolutional local eligibility matches exact BP when instructed."""
import math
import os
import sys

import torch
import torch.nn.functional as F

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from sdil.conv import (CIFARHierarchicalFAResNet, CIFARKPMixedTrafficResNet,
                       CIFARKPResNet, CIFARLocalResNet, CIFARSDILResNet,
                       ConvSDILConfig,
                       channel_subspace_apical_calibration,
                       conv_hierarchical_step, conv_kolen_pollack_step,
                       conv_kp_mixed_traffic_step,
                       conv_local_step,
                       hierarchical_mirror_observations,
                       hierarchical_parameter_subspace_calibration,
                       normalized_residual_mirror_update,
                       normalized_response_mirror_update,
                       simultaneous_conv_node_perturbation,
                       vectorizer_subspace_apical_calibration)


def architecture_checks():
    expected = {
        8: (7, 74810),
        20: (19, 268346),
        32: (31, 461882),
        56: (55, 848954),
    }
    for depth, (hidden, parameters) in expected.items():
        net = CIFARLocalResNet(depth=depth)
        assert net.n_hidden == hidden
        assert net.n_forward_parameters == parameters
        assert len(net.blocks) == 3 * ((depth - 2) // 6)
        assert len(net.W) + 1 == depth
    batchnorm_parameters = {8: 75290, 20: 269722, 32: 464154, 56: 853018}
    for depth, parameters in batchnorm_parameters.items():
        net = CIFARLocalResNet(
            depth=depth, normalization="batchnorm", residual_scale=1.0)
        assert net.n_forward_parameters == parameters
    for depth in (7, 9, 21):
        try:
            CIFARLocalResNet(depth=depth)
        except ValueError:
            pass
        else:
            raise AssertionError(f"invalid depth {depth} was accepted")

    x = torch.arange(2 * 3 * 8 * 8, dtype=torch.float32).reshape(2, 3, 8, 8)
    shortcut = CIFARLocalResNet._option_a_shortcut(x, 6, 2)
    assert tuple(shortcut.shape) == (2, 6, 4, 4)
    assert torch.count_nonzero(shortcut[:, 0]) == 0
    assert torch.count_nonzero(shortcut[:, -2:]) == 0
    assert torch.equal(shortcut[:, 1:4], x[:, :, ::2, ::2])


def exact_local_gradient_check():
    torch.manual_seed(123)
    batch = 3
    x = torch.randn(batch, 3, 32, 32)
    y = torch.tensor([0, 3, 8])
    local = CIFARLocalResNet(depth=8, base_width=4, seed=19)
    bp = CIFARLocalResNet(depth=8, base_width=4, seed=19)

    parameters = local.W + [local.W_out, local.b_out]
    for parameter in parameters:
        parameter.requires_grad_(True)
    forward = local.forward(x, return_cache=True)
    for hidden in forward["hiddens"]:
        hidden.retain_grad()
    loss = F.cross_entropy(forward["logits"], y)
    loss.backward()

    # .backward() differentiated a batch-mean loss.  Multiplying hidden grads
    # by B recovers the per-example convention consumed by the local rule.
    teaching = [-batch * hidden.grad for hidden in forward["hiddens"]]
    output_error = (torch.softmax(forward["logits"].detach(), dim=1)
                    - F.one_hot(y, local.n_classes))
    (directions, gamma_directions, beta_directions,
     out_direction, bias_direction) = local.local_ascent_directions(
        teaching, output_error, forward)
    assert gamma_directions == beta_directions == []

    relative_errors = []
    for direction, parameter in zip(directions, local.W):
        absolute = (direction + parameter.grad).abs().max()
        scale = parameter.grad.abs().max().clamp_min(1e-12)
        relative_errors.append(float(absolute / scale))
    output_abs = float((out_direction + local.W_out.grad).abs().max())
    bias_abs = float((bias_direction + local.b_out.grad).abs().max())
    assert max(relative_errors) < 3e-5
    assert output_abs < 2e-6 and bias_abs < 2e-6

    for parameter in parameters:
        parameter.requires_grad_(False)
    eta = 0.017
    local.apply_ascent(directions, out_direction, bias_direction, eta)
    bp_loss = bp.bp_step(x, y, eta)
    assert abs(float(loss.detach()) - bp_loss) < 1e-7
    parameter_differences = [
        float((left - right).abs().max())
        for left, right in zip(
            local.W + [local.W_out, local.b_out],
            bp.W + [bp.W_out, bp.b_out])]
    assert max(parameter_differences) < 2e-7
    return {
        "max_relative_local_gradient_error": max(relative_errors),
        "output_absolute_error": output_abs,
        "post_update_parameter_max_error": max(parameter_differences),
    }


def exact_batchnorm_local_gradient_check():
    torch.manual_seed(31)
    batch = 4
    x = torch.randn(batch, 3, 32, 32)
    y = torch.tensor([0, 1, 2, 3])
    common = dict(
        depth=8, base_width=2, seed=29,
        normalization="batchnorm", residual_scale=1.0)
    local = CIFARLocalResNet(**common)
    bp = CIFARLocalResNet(**common)
    parameters = local.W + local.gamma + local.beta + [local.W_out, local.b_out]
    for parameter in parameters:
        parameter.requires_grad_(True)
    forward = local.forward(
        x, return_cache=True, training=True, update_stats=True)
    for hidden in forward["hiddens"]:
        hidden.retain_grad()
    loss = F.cross_entropy(forward["logits"], y)
    loss.backward()
    teaching = [-batch * hidden.grad for hidden in forward["hiddens"]]
    output_error = (torch.softmax(forward["logits"].detach(), dim=1)
                    - F.one_hot(y, 10))
    (directions, gamma_directions, beta_directions,
     out_direction, bias_direction) = local.local_ascent_directions(
        teaching, output_error, forward)
    groups = (
        (directions, local.W),
        (gamma_directions, local.gamma),
        (beta_directions, local.beta),
    )
    relative_errors = []
    for direction_group, parameter_group in groups:
        for direction, parameter in zip(direction_group, parameter_group):
            absolute = (direction + parameter.grad).abs().max()
            relative_errors.append(float(
                absolute / parameter.grad.abs().max().clamp_min(1e-12)))
    assert max(relative_errors) < 3e-5
    for parameter in parameters:
        parameter.requires_grad_(False)

    eta = 0.013
    local.apply_ascent(
        directions, out_direction, bias_direction, eta,
        gamma_directions=gamma_directions, beta_directions=beta_directions)
    bp.bp_step(x, y, eta)
    parameter_differences = [
        float((left - right).abs().max())
        for left, right in zip(
            local.W + local.gamma + local.beta + [local.W_out, local.b_out],
            bp.W + bp.gamma + bp.beta + [bp.W_out, bp.b_out])]
    running_differences = [
        float((left - right).abs().max())
        for left, right in zip(
            local.running_mean + local.running_var,
            bp.running_mean + bp.running_var)]
    assert max(parameter_differences) < 2e-7
    assert max(running_differences) == 0.0

    running_before = [value.clone() for value in local.running_mean + local.running_var]
    clean = local.forward(x, training=True, update_stats=False)
    simultaneous_conv_node_perturbation(
        local, x, y, clean, sigma=1e-3, n_directions=1,
        generator=torch.Generator(device="cpu").manual_seed(9))
    assert all(torch.equal(before, after) for before, after in zip(
        running_before, local.running_mean + local.running_var))
    assert torch.equal(local.logits(x), local.logits(x))
    return {
        "batchnorm_max_relative_local_gradient_error": max(relative_errors),
        "batchnorm_post_update_parameter_max_error": max(parameter_differences),
    }


def perturbation_checks():
    net = CIFARLocalResNet(depth=8, base_width=4, seed=3)
    x = torch.randn(2, 3, 32, 32)
    clean = net.forward(x)
    perturbations = [torch.zeros_like(hidden) for hidden in clean["hiddens"]]
    perturbed = net.forward(x, perturbations=perturbations)
    assert torch.equal(clean["logits"], perturbed["logits"])
    perturbations[0] = torch.ones_like(perturbations[0]) * 0.01
    changed = net.forward(x, perturbations=perturbations)
    assert not torch.equal(clean["logits"], changed["logits"])
    try:
        net.forward(x, perturbations=perturbations[:-1])
    except ValueError:
        pass
    else:
        raise AssertionError("short perturbation list was accepted")


def perturbation_estimator_check():
    """Antithetic finite differences equal the simultaneous hidden JVP."""
    torch.manual_seed(3)
    batch = 2
    net = CIFARSDILResNet(
        depth=8, base_width=2, seed=4, dtype=torch.float64)
    x = torch.randn(batch, 3, 32, 32, dtype=torch.float64)
    y = torch.tensor([2, 7])
    parameters = net.W + [net.W_out, net.b_out]
    for parameter in parameters:
        parameter.requires_grad_(True)
    clean = net.forward(x, return_cache=True)
    for hidden in clean["hiddens"]:
        hidden.retain_grad()
    F.cross_entropy(clean["logits"], y).backward()
    generator = torch.Generator(device="cpu").manual_seed(99)
    targets, diagnostics = simultaneous_conv_node_perturbation(
        net, x, y, clean, sigma=1e-6, n_directions=1,
        generator=generator, return_diagnostics=True)
    directions = diagnostics["directions"][0]
    derivative = diagnostics["directional_derivatives"][0]
    assert derivative["coupling"] == "per_example_objective"
    finite_difference = derivative["scaled_directional"]
    exact = sum(
        (batch * hidden.grad * direction).flatten(1).sum(dim=1)
        for hidden, direction in zip(clean["hiddens"], directions))
    relative = (finite_difference - exact).abs() / exact.abs().clamp_min(1e-12)
    assert float(relative.max()) < 2e-6
    for target, direction in zip(targets, directions):
        expected = -finite_difference[:, None, None, None] * direction
        assert torch.equal(target, expected)
    for parameter in parameters:
        parameter.requires_grad_(False)
    no_norm_relative = float(relative.max())

    batch = 3
    batchnorm = CIFARSDILResNet(
        depth=8, base_width=2, seed=14, dtype=torch.float64,
        normalization="batchnorm", residual_scale=1.0)
    xb = torch.randn(batch, 3, 32, 32, dtype=torch.float64)
    yb = torch.tensor([1, 4, 9])
    parameters = (batchnorm.W + batchnorm.gamma + batchnorm.beta
                  + [batchnorm.W_out, batchnorm.b_out])
    for parameter in parameters:
        parameter.requires_grad_(True)
    clean = batchnorm.forward(xb, training=True, update_stats=False)
    for hidden in clean["hiddens"]:
        hidden.retain_grad()
    F.cross_entropy(clean["logits"], yb).backward()
    targets, diagnostics = simultaneous_conv_node_perturbation(
        batchnorm, xb, yb, clean, sigma=1e-6, n_directions=1,
        generator=torch.Generator(device="cpu").manual_seed(101),
        return_diagnostics=True)
    directions = diagnostics["directions"][0]
    derivative = diagnostics["directional_derivatives"][0]
    assert derivative["coupling"] == "batch_objective"
    scaled = derivative["scaled_directional"]
    exact_sum_directional = sum(
        float((batch * hidden.grad * direction).sum())
        for hidden, direction in zip(clean["hiddens"], directions))
    batchnorm_relative = abs(float(scaled[0]) - exact_sum_directional) / max(
        abs(exact_sum_directional), 1e-12)
    assert batchnorm_relative < 2e-6
    assert torch.equal(scaled, scaled[:1].expand_as(scaled))
    for target, direction in zip(targets, directions):
        expected = -scaled[:, None, None, None] * direction
        assert torch.equal(target, expected)
    for parameter in parameters:
        parameter.requires_grad_(False)
    return {
        "perturbation_jvp_max_relative_error": no_norm_relative,
        "batchnorm_batch_objective_jvp_relative_error": batchnorm_relative,
    }


def channel_subspace_estimator_check():
    """The structured estimator targets representable base/gate moments."""
    torch.manual_seed(71)
    batch = 3
    hiddens = [
        torch.randn(batch, 2, 4, 4, dtype=torch.float64),
        torch.randn(batch, 3, 2, 2, dtype=torch.float64),
    ]
    negative_gradients = [torch.randn_like(value) for value in hiddens]
    estimated_base = [torch.zeros(
        batch, value.shape[1], dtype=value.dtype) for value in hiddens]
    estimated_gate = [torch.zeros_like(value) for value in estimated_base]
    generator = torch.Generator(device="cpu").manual_seed(211)
    directions = 4096
    inverse_sqrt_two = 1.0 / (2.0 ** 0.5)
    for _ in range(directions):
        hidden_directions = []
        base_random = []
        gate_random = []
        for hidden in hiddens:
            shape = (batch, hidden.shape[1])
            base = torch.empty(shape, dtype=hidden.dtype).bernoulli_(
                0.5, generator=generator).mul_(2).sub_(1)
            gate = torch.empty_like(base).bernoulli_(
                0.5, generator=generator).mul_(2).sub_(1)
            hidden_directions.append((
                base[:, :, None, None]
                + torch.tanh(hidden) * gate[:, :, None, None])
                * inverse_sqrt_two)
            base_random.append(base)
            gate_random.append(gate)
        # The exact loss derivative uses g=-negative_gradient and remains
        # per-example without BatchNorm coupling.
        directional = -sum(
            (gradient * direction).flatten(1).sum(dim=1)
            for gradient, direction in zip(
                negative_gradients, hidden_directions))
        for index, (hidden, base, gate) in enumerate(zip(
                hiddens, base_random, gate_random)):
            spatial = hidden.shape[2] * hidden.shape[3]
            scale = -(2.0 ** 0.5) / (spatial * directions)
            estimated_base[index].add_(directional[:, None] * base, alpha=scale)
            estimated_gate[index].add_(directional[:, None] * gate, alpha=scale)
    exact_base = [value.mean(dim=(2, 3)) for value in negative_gradients]
    exact_gate = [(value * torch.tanh(hidden)).mean(dim=(2, 3))
                  for value, hidden in zip(negative_gradients, hiddens)]
    estimated = torch.cat([
        value.flatten() for pair in zip(estimated_base, estimated_gate)
        for value in pair])
    exact = torch.cat([
        value.flatten() for pair in zip(exact_base, exact_gate)
        for value in pair])
    cosine = float(F.cosine_similarity(estimated, exact, dim=0))
    norm_ratio = float(estimated.norm() / exact.norm())
    assert cosine > 0.985
    assert 0.90 < norm_ratio < 1.10

    # The executable antithetic implementation must match the same structured
    # directional derivative, not an autograd surrogate.
    # Use a fixed nondegenerate point.  Width-one, zero-bias ReLU networks can
    # contain structurally exact-zero preactivations, where central differences
    # and PyTorch's chosen subgradient need not agree even as sigma -> 0.
    torch.manual_seed(3)
    net = CIFARSDILResNet(
        depth=8, base_width=2, seed=72, dtype=torch.float64,
        vectorizer_mode="channel_gated")
    x = torch.randn(2, 3, 32, 32, dtype=torch.float64)
    y = torch.tensor([2, 8])
    parameters = net.W + [net.W_out, net.b_out]
    for parameter in parameters:
        parameter.requires_grad_(True)
    clean = net.forward(x, return_cache=True)
    for hidden in clean["hiddens"]:
        hidden.retain_grad()
    loss = F.cross_entropy(clean["logits"], y)
    loss.backward()
    output_signal = (torch.softmax(clean["logits"].detach(), dim=1)
                     - F.one_hot(y, 10).to(torch.float64))
    _, diagnostics = channel_subspace_apical_calibration(
        net, x, y, clean, output_signal, sigma=1e-6,
        n_directions=1, eta=0.0,
        generator=torch.Generator(device="cpu").manual_seed(307),
        return_diagnostics=True)
    hidden_direction = diagnostics["directions"][0]["hidden"]
    finite_difference = diagnostics["directional_derivatives"][0][
        "scaled_directional"]
    exact_directional = sum(
        (x.shape[0] * hidden.grad * direction).flatten(1).sum(dim=1)
        for hidden, direction in zip(clean["hiddens"], hidden_direction))
    relative = ((finite_difference - exact_directional).abs()
                / exact_directional.abs().clamp_min(1e-12))
    assert float(relative.max()) < 2e-6
    for parameter in parameters:
        parameter.requires_grad_(False)

    # The local A/G update must equal the full-field delta rule after replacing
    # only its two target moments with the structured causal estimates.
    eta = 0.0023
    before_a = [value.clone() for value in net.A]
    before_g = [value.clone() for value in net.A_gate]
    expected_a = []
    expected_g = []
    for index, hidden in enumerate(clean["hiddens"]):
        base = output_signal @ before_a[index].t()
        gate_coefficient = output_signal @ before_g[index].t()
        gate = torch.tanh(hidden.detach())
        mean = gate.mean(dim=(2, 3))
        second = gate.square().mean(dim=(2, 3))
        base_error = diagnostics["target_base"][index] - (
            base + mean * gate_coefficient)
        gate_error = diagnostics["target_gate"][index] - (
            mean * base + second * gate_coefficient)
        expected_a.append(before_a[index] + eta * (
            base_error.t() @ output_signal / x.shape[0]))
        expected_g.append(before_g[index] + eta * (
            gate_error.t() @ output_signal / x.shape[0]))
    channel_subspace_apical_calibration(
        net, x, y, clean, output_signal, sigma=1e-6,
        n_directions=1, eta=eta,
        generator=torch.Generator(device="cpu").manual_seed(307))
    update_error = max(float((actual - expected).abs().max())
                       for actual, expected in zip(
                           net.A + net.A_gate, expected_a + expected_g))
    assert update_error < 1e-14

    return {
        "channel_subspace_moment_cosine": cosine,
        "channel_subspace_moment_norm_ratio": norm_ratio,
        "channel_subspace_jvp_relative_error": float(relative.max()),
        "channel_subspace_delta_rule_absolute_error": update_error,
    }


def vectorizer_subspace_estimator_check():
    """Direct A/G perturbations are unbiased and lower variance at batch 128."""
    torch.manual_seed(81)
    batch = 128
    output_dim = 5
    hiddens = [
        torch.randn(batch, 2, 4, 4, dtype=torch.float64),
        torch.randn(batch, 3, 2, 2, dtype=torch.float64),
    ]
    negative_gradients = [torch.randn_like(value) for value in hiddens]
    output_signal = torch.randn(batch, output_dim, dtype=torch.float64)
    exact_pairs = []
    for hidden, target in zip(hiddens, negative_gradients):
        exact_pairs.extend([
            target.mean(dim=(2, 3)).t() @ output_signal / batch,
            (target * torch.tanh(hidden)).mean(dim=(2, 3)).t()
            @ output_signal / batch,
        ])
    exact = torch.cat([value.flatten() for value in exact_pairs])
    coefficient_sum = torch.zeros_like(exact)
    vectorizer_sum = torch.zeros_like(exact)
    coefficient_mse = 0.0
    vectorizer_mse = 0.0
    directions = 2048
    generator = torch.Generator(device="cpu").manual_seed(401)
    inverse_sqrt_two = 1.0 / (2.0 ** 0.5)
    for _ in range(directions):
        random_coefficients = []
        hidden_directions = []
        for hidden in hiddens:
            shape = (batch, hidden.shape[1])
            base = torch.empty(shape, dtype=hidden.dtype).bernoulli_(
                0.5, generator=generator).mul_(2).sub_(1)
            gate = torch.empty_like(base).bernoulli_(
                0.5, generator=generator).mul_(2).sub_(1)
            random_coefficients.append((base, gate))
            hidden_directions.append((
                base[:, :, None, None]
                + torch.tanh(hidden) * gate[:, :, None, None])
                * inverse_sqrt_two)
        directional = -sum((target * direction).sum()
                           for target, direction in zip(
                               negative_gradients, hidden_directions))
        coefficient_values = []
        for hidden, (base, gate) in zip(hiddens, random_coefficients):
            spatial = hidden.shape[2] * hidden.shape[3]
            base_target = -(2.0 ** 0.5) * directional * base / spatial
            gate_target = -(2.0 ** 0.5) * directional * gate / spatial
            coefficient_values.extend([
                base_target.t() @ output_signal / batch,
                gate_target.t() @ output_signal / batch,
            ])
        coefficient_sample = torch.cat(
            [value.flatten() for value in coefficient_values])

        random_matrices = []
        hidden_directions = []
        for hidden in hiddens:
            shape = (hidden.shape[1], output_dim)
            base = torch.empty(shape, dtype=hidden.dtype).bernoulli_(
                0.5, generator=generator).mul_(2).sub_(1)
            gate = torch.empty_like(base).bernoulli_(
                0.5, generator=generator).mul_(2).sub_(1)
            base_field = output_signal @ base.t()
            gate_field = output_signal @ gate.t()
            random_matrices.append((base, gate))
            hidden_directions.append((
                base_field[:, :, None, None]
                + torch.tanh(hidden) * gate_field[:, :, None, None])
                * inverse_sqrt_two)
        directional = -sum((target * direction).sum()
                           for target, direction in zip(
                               negative_gradients, hidden_directions))
        vectorizer_values = []
        for hidden, (base, gate) in zip(hiddens, random_matrices):
            spatial = hidden.shape[2] * hidden.shape[3]
            scale = -(2.0 ** 0.5) * directional / (batch * spatial)
            vectorizer_values.extend([scale * base, scale * gate])
        vectorizer_sample = torch.cat(
            [value.flatten() for value in vectorizer_values])
        coefficient_sum.add_(coefficient_sample)
        vectorizer_sum.add_(vectorizer_sample)
        coefficient_mse += float((coefficient_sample - exact).square().mean())
        vectorizer_mse += float((vectorizer_sample - exact).square().mean())
    coefficient_mean = coefficient_sum / directions
    vectorizer_mean = vectorizer_sum / directions
    vectorizer_cosine = float(F.cosine_similarity(
        vectorizer_mean, exact, dim=0))
    vectorizer_norm_ratio = float(vectorizer_mean.norm() / exact.norm())
    variance_ratio = vectorizer_mse / coefficient_mse
    assert vectorizer_cosine > 0.95
    assert 0.90 < vectorizer_norm_ratio < 1.10
    assert variance_ratio < 0.25

    # Match the executable forward-only derivative and its exact A/G update.
    torch.manual_seed(3)
    net = CIFARSDILResNet(
        depth=8, base_width=2, seed=82, dtype=torch.float64,
        vectorizer_mode="channel_gated")
    x = torch.randn(2, 3, 32, 32, dtype=torch.float64)
    y = torch.tensor([1, 6])
    parameters = net.W + [net.W_out, net.b_out]
    for parameter in parameters:
        parameter.requires_grad_(True)
    clean = net.forward(x, return_cache=True)
    for hidden in clean["hiddens"]:
        hidden.retain_grad()
    F.cross_entropy(clean["logits"], y).backward()
    output_error = (torch.softmax(clean["logits"].detach(), dim=1)
                    - F.one_hot(y, 10).to(torch.float64))
    _, diagnostics = vectorizer_subspace_apical_calibration(
        net, x, y, clean, output_error, sigma=1e-6,
        n_directions=1, eta=0.0,
        generator=torch.Generator(device="cpu").manual_seed(503),
        return_diagnostics=True)
    finite_difference = diagnostics["directional_derivatives"][0][
        "scaled_directional"]
    exact_directional = sum(
        (x.shape[0] * hidden.grad * direction).sum()
        for hidden, direction in zip(
            clean["hiddens"], diagnostics["directions"][0]["hidden"]))
    jvp_relative = float((finite_difference - exact_directional).abs()
                         / exact_directional.abs().clamp_min(1e-12))
    assert jvp_relative < 2e-6
    for parameter in parameters:
        parameter.requires_grad_(False)

    eta = 0.0017
    before_a = [value.clone() for value in net.A]
    before_g = [value.clone() for value in net.A_gate]
    expected_a = []
    expected_g = []
    for index, hidden in enumerate(clean["hiddens"]):
        base = output_error @ before_a[index].t()
        gate_coefficient = output_error @ before_g[index].t()
        gate = torch.tanh(hidden.detach())
        mean = gate.mean(dim=(2, 3))
        second = gate.square().mean(dim=(2, 3))
        base_prediction = (base + mean * gate_coefficient).t() @ output_error / 2
        gate_prediction = (
            mean * base + second * gate_coefficient).t() @ output_error / 2
        expected_a.append(before_a[index] + eta * (
            diagnostics["target_base"][index] - base_prediction))
        expected_g.append(before_g[index] + eta * (
            diagnostics["target_gate"][index] - gate_prediction))
    vectorizer_subspace_apical_calibration(
        net, x, y, clean, output_error, sigma=1e-6,
        n_directions=1, eta=eta,
        generator=torch.Generator(device="cpu").manual_seed(503))
    update_error = max(float((actual - expected).abs().max())
                       for actual, expected in zip(
                           net.A + net.A_gate, expected_a + expected_g))
    assert update_error < 1e-14

    batchnorm = CIFARSDILResNet(
        depth=8, base_width=2, seed=83, dtype=torch.float64,
        normalization="batchnorm", residual_scale=1.0,
        vectorizer_mode="channel_gated")
    xb = torch.randn(3, 3, 32, 32, dtype=torch.float64)
    yb = torch.tensor([0, 4, 9])
    parameters = (batchnorm.W + batchnorm.gamma + batchnorm.beta
                  + [batchnorm.W_out, batchnorm.b_out])
    for parameter in parameters:
        parameter.requires_grad_(True)
    clean_b = batchnorm.forward(xb, training=True, update_stats=False)
    for hidden in clean_b["hiddens"]:
        hidden.retain_grad()
    F.cross_entropy(clean_b["logits"], yb).backward()
    output_b = (torch.softmax(clean_b["logits"].detach(), dim=1)
                - F.one_hot(yb, 10).to(torch.float64))
    _, diagnostics_b = vectorizer_subspace_apical_calibration(
        batchnorm, xb, yb, clean_b, output_b, sigma=1e-6,
        n_directions=1, eta=0.0,
        generator=torch.Generator(device="cpu").manual_seed(509),
        return_diagnostics=True)
    finite_b = diagnostics_b["directional_derivatives"][0][
        "scaled_directional"]
    exact_b = sum(
        (xb.shape[0] * hidden.grad * direction).sum()
        for hidden, direction in zip(
            clean_b["hiddens"], diagnostics_b["directions"][0]["hidden"]))
    batchnorm_relative = float(
        (finite_b - exact_b).abs() / exact_b.abs().clamp_min(1e-12))
    assert batchnorm_relative < 2e-6
    for parameter in parameters:
        parameter.requires_grad_(False)
    return {
        "vectorizer_subspace_mean_cosine": vectorizer_cosine,
        "vectorizer_subspace_mean_norm_ratio": vectorizer_norm_ratio,
        "vectorizer_vs_coefficient_mse_ratio": variance_ratio,
        "vectorizer_subspace_jvp_relative_error": jvp_relative,
        "vectorizer_subspace_batchnorm_jvp_relative_error": batchnorm_relative,
        "vectorizer_subspace_delta_rule_absolute_error": update_error,
    }


def hierarchical_feedback_checks():
    """The residual feedback graph becomes exact only under an audit copy."""
    torch.manual_seed(91)
    exact = CIFARHierarchicalFAResNet(
        depth=8, base_width=2, seed=92, dtype=torch.float64,
        normalization="batchnorm", residual_scale=1.0)
    exact.Q = [value.clone() for value in exact.W]
    exact.R_out.copy_(-exact.W_out.t())
    x = torch.randn(3, 3, 32, 32, dtype=torch.float64)
    y = torch.tensor([1, 5, 8])
    parameters = (exact.W + exact.gamma + exact.beta
                  + [exact.W_out, exact.b_out])
    for parameter in parameters:
        parameter.requires_grad_(True)
    forward = exact.forward(x, return_cache=True, training=True,
                            update_stats=False)
    gradients = torch.autograd.grad(
        F.cross_entropy(forward["logits"], y), forward["hiddens"])
    output_error = (torch.softmax(forward["logits"].detach(), dim=1)
                    - F.one_hot(y, 10).to(torch.float64))
    teaching = exact.hierarchical_teaching(output_error, forward)
    relative = [float((signal + x.shape[0] * gradient).abs().max()
                      / gradient.abs().max().clamp_min(1e-30)
                      / x.shape[0])
                for signal, gradient in zip(teaching, gradients)]
    assert max(relative) < 2e-12
    for parameter in parameters:
        parameter.requires_grad_(False)

    # With independent feedback the forward model is bitwise unchanged, while
    # changing only the feedback seed changes Q/R. This is the actual baseline.
    left = CIFARHierarchicalFAResNet(
        depth=8, base_width=2, seed=93, feedback_seed=1001)
    right = CIFARHierarchicalFAResNet(
        depth=8, base_width=2, seed=93, feedback_seed=1002)
    assert all(torch.equal(a, b) for a, b in zip(
        left.W + [left.W_out], right.W + [right.W_out]))
    assert any(not torch.equal(a, b) for a, b in zip(
        left.Q + [left.R_out], right.Q + [right.R_out]))

    # The local update must match exact BP when feedback is explicitly copied
    # in this audit-only comparator. Actual HFA never performs this copy.
    local = CIFARHierarchicalFAResNet(
        depth=8, base_width=2, seed=94, normalization="batchnorm",
        residual_scale=1.0)
    bp = CIFARLocalResNet(
        depth=8, base_width=2, seed=94, normalization="batchnorm",
        residual_scale=1.0)
    local.Q = [value.clone() for value in local.W]
    local.R_out.copy_(-local.W_out.t())
    xf = torch.randn(4, 3, 32, 32)
    yf = torch.tensor([0, 2, 5, 9])
    eta = 0.013
    conv_hierarchical_step(
        local, xf, yf, ConvSDILConfig(
            eta=eta, eta_output=eta, eta_A=0.0, momentum=0.0,
            weight_decay=0.0, learn_A=False))
    bp.bp_step(xf, yf, eta, momentum=0.0, weight_decay=0.0)
    parameter_error = max(float((a - b).abs().max()) for a, b in zip(
        local.W + local.gamma + local.beta + [local.W_out, local.b_out],
        bp.W + bp.gamma + bp.beta + [bp.W_out, bp.b_out]))
    running_error = max(float((a - b).abs().max()) for a, b in zip(
        local.running_mean + local.running_var,
        bp.running_mean + bp.running_var))
    assert parameter_error < 2e-7
    assert running_error == 0.0
    return {
        "hierarchical_symmetric_hidden_relative_error": max(relative),
        "hierarchical_symmetric_update_absolute_error": parameter_error,
        "hierarchical_feedback_to_forward_mac_ratio": (
            local.apical_macs_per_example / local.forward_macs_per_example),
    }


def hierarchical_parameter_calibration_checks():
    """Audit the causal JVP and the exact local Q/R delta-rule moments."""
    torch.manual_seed(109)
    net = CIFARHierarchicalFAResNet(
        depth=8, base_width=2, seed=110, dtype=torch.float64,
        normalization="batchnorm", residual_scale=1.0)
    x = torch.randn(3, 3, 32, 32, dtype=torch.float64)
    y = torch.tensor([0, 4, 7])
    parameters = net.W + net.gamma + net.beta + [net.W_out, net.b_out]
    for parameter in parameters:
        parameter.requires_grad_(True)
    forward = net.forward(
        x, return_cache=True, training=True, update_stats=False)
    loss = F.cross_entropy(forward["logits"], y)
    hidden_gradients = torch.autograd.grad(loss, forward["hiddens"])
    output_signal = (torch.softmax(forward["logits"].detach(), dim=1)
                     - F.one_hot(y, 10).to(torch.float64))
    _, diagnostic = hierarchical_parameter_subspace_calibration(
        net, x, y, forward, output_signal, sigma=1e-5,
        n_directions=1, eta=0.0,
        generator=torch.Generator().manual_seed(111),
        return_diagnostics=True)
    directions = diagnostic["directions"][0]["hidden"]
    exact_directional = x.shape[0] * sum(
        (gradient * direction).sum()
        for gradient, direction in zip(hidden_gradients, directions))
    estimated_directional = diagnostic[
        "directional_derivatives"][0]["scaled_directional"]
    jvp_relative = float((estimated_directional - exact_directional).abs()
                         / exact_directional.abs().clamp_min(1e-30))
    assert jvp_relative < 3e-7
    for parameter in parameters:
        parameter.requires_grad_(False)

    # Under an audit-only symmetric copy, the hierarchical field is the exact
    # negative gradient. Consequently every local Q/R predicted moment equals
    # its exact causal regression target, including option-A shortcut terms.
    exact = CIFARHierarchicalFAResNet(
        depth=8, base_width=2, seed=112, dtype=torch.float64,
        normalization="batchnorm", residual_scale=1.0)
    exact.Q = [value.clone() for value in exact.W]
    exact.R_out.copy_(-exact.W_out.t())
    for parameter in exact.W + exact.gamma + exact.beta + [
            exact.W_out, exact.b_out]:
        parameter.requires_grad_(True)
    clean = exact.forward(
        x, return_cache=True, training=True, update_stats=False)
    gradients = torch.autograd.grad(
        F.cross_entropy(clean["logits"], y), clean["hiddens"])
    negative = [-x.shape[0] * value.detach() for value in gradients]
    signal = (torch.softmax(clean["logits"].detach(), dim=1)
              - F.one_hot(y, 10).to(torch.float64))
    teaching, contexts, recipients = exact.hierarchical_teaching(
        signal, clean, return_edge_contexts=True)
    numerator = 0.0
    denominator = 0.0
    for index in range(1, len(exact.Q)):
        recipient = recipients[index]
        spec = exact.layer_specs[index]
        spatial = (negative[recipient].shape[2]
                   * negative[recipient].shape[3])
        target = torch.nn.grad.conv2d_weight(
            negative[recipient], exact.Q[index].shape, contexts[index],
            stride=spec.stride, padding=spec.padding) / (x.shape[0] * spatial)
        prediction = torch.nn.grad.conv2d_weight(
            teaching[recipient], exact.Q[index].shape, contexts[index],
            stride=spec.stride, padding=spec.padding) / (x.shape[0] * spatial)
        numerator += float((target - prediction).square().sum())
        denominator += float(target.square().sum())
    target_r = negative[-1].mean(dim=(2, 3)).t() @ signal / x.shape[0]
    prediction_r = teaching[-1].mean(dim=(2, 3)).t() @ signal / x.shape[0]
    numerator += float((target_r - prediction_r).square().sum())
    denominator += float(target_r.square().sum())
    delta_rule_relative = math.sqrt(numerator / max(denominator, 1e-300))
    assert delta_rule_relative < 2e-12
    return {
        "hierarchical_parameter_subspace_jvp_relative_error": jvp_relative,
        "hierarchical_parameter_delta_rule_relative_error": delta_rule_relative,
    }


def normalized_response_mirror_checks():
    """Audit local response estimation and absence of W access in the update."""
    net = CIFARHierarchicalFAResNet(
        depth=8, base_width=2, seed=121, dtype=torch.float64,
        normalization="batchnorm")
    observations = hierarchical_mirror_observations(
        net, batch_size=16, noise_std=1.0,
        generator=torch.Generator().manual_seed(122))
    metrics, _ = normalized_response_mirror_update(
        net, observations, eta=1.0)
    pairs = list(zip(net.Q[1:], net.W[1:])) + [
        (net.R_out, -net.W_out.t())]
    cosines = [float(F.cosine_similarity(
        feedback.flatten(), target.flatten(), dim=0))
        for feedback, target in pairs]
    norm_ratios = [float(feedback.norm() / target.norm())
                   for feedback, target in pairs]
    assert sum(cosines) / len(cosines) > 0.985
    assert min(cosines) > 0.95
    assert min(norm_ratios) > 0.90 and max(norm_ratios) < 1.10

    # The update consumes observations only. Changing every forward parameter
    # after those observations were generated must not change the Q/R update.
    left = CIFARHierarchicalFAResNet(
        depth=8, base_width=2, seed=123, dtype=torch.float64)
    right = CIFARHierarchicalFAResNet(
        depth=8, base_width=2, seed=123, dtype=torch.float64)
    shared_observations = hierarchical_mirror_observations(
        left, batch_size=2, generator=torch.Generator().manual_seed(124))
    for value in right.W + [right.W_out]:
        value.normal_(generator=torch.Generator().manual_seed(value.numel()))
    normalized_response_mirror_update(left, shared_observations, eta=0.2)
    normalized_response_mirror_update(right, shared_observations, eta=0.2)
    independence_error = max(float((a - b).abs().max()) for a, b in zip(
        left.Q[1:] + [left.R_out], right.Q[1:] + [right.R_out]))
    assert independence_error == 0.0

    # Residual-response LMS has a per-observation exact fixed point: its update
    # is zero, not merely zero in expectation, when Q/R match the forward maps.
    fixed = CIFARHierarchicalFAResNet(
        depth=8, base_width=2, seed=125, dtype=torch.float64)
    fixed.Q = [value.clone() for value in fixed.W]
    fixed.R_out.copy_(-fixed.W_out.t())
    fixed_observations = hierarchical_mirror_observations(
        fixed, batch_size=2, generator=torch.Generator().manual_seed(126))
    fixed_metrics = normalized_residual_mirror_update(
        fixed, fixed_observations, eta=1.0)
    assert fixed_metrics["mirror_update_rms"] < 1e-14
    assert fixed_metrics["mirror_response_residual_fraction"] < 1e-14

    residual_left = CIFARHierarchicalFAResNet(
        depth=8, base_width=2, seed=127, dtype=torch.float64)
    residual_right = CIFARHierarchicalFAResNet(
        depth=8, base_width=2, seed=127, dtype=torch.float64)
    residual_observations = hierarchical_mirror_observations(
        residual_left, batch_size=2,
        generator=torch.Generator().manual_seed(128))
    for value in residual_right.W + [residual_right.W_out]:
        value.normal_(generator=torch.Generator().manual_seed(value.numel() + 1))
    normalized_residual_mirror_update(
        residual_left, residual_observations, eta=0.2)
    normalized_residual_mirror_update(
        residual_right, residual_observations, eta=0.2)
    residual_independence_error = max(float((a - b).abs().max())
                                      for a, b in zip(
        residual_left.Q[1:] + [residual_left.R_out],
        residual_right.Q[1:] + [residual_right.R_out]))
    assert residual_independence_error == 0.0
    return {
        "mirror_estimate_mean_forward_cosine": sum(cosines) / len(cosines),
        "mirror_estimate_min_forward_cosine": min(cosines),
        "mirror_estimate_min_norm_ratio": min(norm_ratios),
        "mirror_estimate_max_norm_ratio": max(norm_ratios),
        "mirror_update_forward_parameter_independence_error": independence_error,
        "mirror_update_rms": metrics["mirror_update_rms"],
        "residual_mirror_exact_fixed_point_update_rms": fixed_metrics[
            "mirror_update_rms"],
        "residual_mirror_exact_fixed_point_fraction": fixed_metrics[
            "mirror_response_residual_fraction"],
        "residual_mirror_forward_parameter_independence_error": (
            residual_independence_error),
    }


def kolen_pollack_checks():
    """KP's reciprocal correlations are local and preserve exact symmetry."""
    torch.manual_seed(127)
    common = dict(
        depth=8, base_width=2, seed=53, dtype=torch.float64,
        normalization="batchnorm", residual_scale=1.0)
    net = CIFARKPResNet(**common)
    for index in range(1, len(net.Q)):
        net.Q[index].copy_(net.W[index])
    net.R_out.copy_(-net.W_out.t())
    x = torch.randn(3, 3, 32, 32, dtype=torch.float64)
    y = torch.tensor([1, 4, 7])
    forward = net.forward(
        x, return_cache=True, training=True, update_stats=False)
    output_error = (torch.softmax(forward["logits"], dim=1)
                    - F.one_hot(y, 10).to(torch.float64))
    teaching = net.hierarchical_teaching(output_error, forward)
    (forward_directions, gamma_directions, beta_directions,
     output_weight, output_bias) = net.local_ascent_directions(
        teaching, output_error, forward)
    reciprocal_directions, reciprocal_readout = (
        net.reciprocal_feedback_directions(
            teaching, output_error, forward))
    direction_error = max([
        float((left - right).abs().max())
        for left, right in zip(forward_directions[1:], reciprocal_directions[1:])
    ] + [float((reciprocal_readout + output_weight.t()).abs().max())])
    assert direction_error < 1e-14

    # Once local activities have been observed, neither forward nor feedback
    # parameter values may alter the independently formed reciprocal update.
    before = [value.clone() for value in reciprocal_directions[1:]]
    before_readout = reciprocal_readout.clone()
    for value in net.W + net.Q:
        value.add_(torch.randn_like(value))
    net.W_out.add_(torch.randn_like(net.W_out))
    net.R_out.add_(torch.randn_like(net.R_out))
    independent_directions, independent_readout = (
        net.reciprocal_feedback_directions(
            teaching, output_error, forward))
    independence_error = max([
        float((left - right).abs().max())
        for left, right in zip(before, independent_directions[1:])
    ] + [float((before_readout - independent_readout).abs().max())])
    assert independence_error == 0.0

    # A fresh symmetric state must remain symmetric under two momentum steps.
    net = CIFARKPResNet(**common)
    for index in range(1, len(net.Q)):
        net.Q[index].copy_(net.W[index])
    net.R_out.copy_(-net.W_out.t())
    for _ in range(2):
        forward = net.forward(
            x, return_cache=True, training=True, update_stats=False)
        output_error = (torch.softmax(forward["logits"], dim=1)
                        - F.one_hot(y, 10).to(torch.float64))
        teaching = net.hierarchical_teaching(output_error, forward)
        (forward_directions, gamma_directions, beta_directions,
         output_weight, output_bias) = net.local_ascent_directions(
            teaching, output_error, forward)
        reciprocal_directions, reciprocal_readout = (
            net.reciprocal_feedback_directions(
                teaching, output_error, forward))
        net.apply_reciprocal_ascent(
            reciprocal_directions, reciprocal_readout,
            eta_hidden=0.013, eta_output=0.017, momentum=0.9,
            weight_decay=1e-4)
        net.apply_ascent(
            forward_directions, output_weight, output_bias,
            eta_hidden=0.013, eta_output=0.017, momentum=0.9,
            weight_decay=1e-4, gamma_directions=gamma_directions,
            beta_directions=beta_directions)
    symmetry_error = max([
        float((net.Q[index] - net.W[index]).abs().max())
        for index in range(1, len(net.Q))
    ] + [float((net.R_out + net.W_out.t()).abs().max())])
    assert symmetry_error < 1e-14

    # Exercise the public training step and ensure it remains graph-free.
    result = conv_kolen_pollack_step(
        net, x, y, ConvSDILConfig(
            eta=1e-3, eta_output=1e-3, momentum=0.9,
            weight_decay=1e-4, learn_A=False, learn_P=False))
    assert math.isfinite(result["loss"])
    assert all(not value.requires_grad for value in
               net.W + net.Q + [net.W_out, net.R_out, net.b_out])
    return {
        "kp_local_direction_absolute_error": direction_error,
        "kp_forward_parameter_independence_error": independence_error,
        "kp_symmetric_update_absolute_error": symmetry_error,
    }


def kp_mixed_traffic_checks():
    """Mixed traffic isolates subtraction from norm and preserves KP locality."""
    torch.manual_seed(211)
    common = dict(
        depth=8, base_width=2, seed=67, dtype=torch.float64,
        normalization="batchnorm", residual_scale=1.0,
        traffic_seed=4000)
    net = CIFARKPMixedTrafficResNet(**common)
    x = torch.randn(5, 3, 32, 32, dtype=torch.float64)
    y = torch.tensor([0, 2, 4, 6, 8])
    forward = net.forward(
        x, return_cache=True, training=True, update_stats=False)
    output_error = (torch.softmax(forward["logits"], dim=1)
                    - F.one_hot(y, 10).to(torch.float64))
    instruction = net.hierarchical_teaching(output_error, forward)

    # Zero traffic and zero predictor collapse all three rules to clean KP.
    zero_errors = []
    for rule in ("raw", "matched", "innovation"):
        components = net.mixed_apical_components(
            instruction, forward["hiddens"], rule)
        zero_errors.extend(float((left - right).abs().max()) for left, right in
                           zip(components["used"], instruction))
    assert max(zero_errors) == 0.0

    calibration = net.calibrate_traffic_gain(
        instruction, forward["hiddens"], target_ratio=4.0)
    ratio_error = max(abs(value - 4.0) for value in
                      calibration["realized_traffic_instruction_rms_ratio"])
    assert ratio_error < 1e-12

    # An exact per-unit predictor removes all predictable traffic.
    for slope, gain, coefficient in zip(
            net.P_traffic, net.traffic_gain, net.B_traffic):
        slope.copy_(gain * coefficient)
    exact = net.mixed_apical_components(
        instruction, forward["hiddens"], "innovation")
    exact_predictor_error = max(float((left - right).abs().max())
                                for left, right in zip(
        exact["innovation"], instruction))
    assert exact_predictor_error < 1e-14

    for slope, bias in zip(net.P_traffic, net.P_traffic_bias):
        slope.zero_()
        bias.zero_()
    closed_form = net.predictor_closed_form_fit(forward["hiddens"])
    fitted = net.mixed_apical_components(
        instruction, forward["hiddens"], "innovation")
    closed_form_error = max(float((left - right).abs().max())
                            for left, right in zip(
        fitted["innovation"], instruction))
    assert closed_form_error < 1e-14
    assert closed_form["residual_traffic_rms_ratio"] < 1e-14
    assert closed_form["max_absolute_residual_soma_slope"] < 1e-14

    for slope, bias in zip(net.P_traffic, net.P_traffic_bias):
        slope.zero_()
        bias.zero_()
    stable_fit = net.predictor_closed_form_fit(
        forward["hiddens"], stability_margin=1e-3)
    assert stable_fit["max_positive_residual_soma_slope"] < 1e-14
    assert stable_fit["min_residual_soma_slope"] < -9e-4
    assert stable_fit["max_applied_stability_margin"] >= 1e-3

    # A deliberately inaccurate slow predictor leaves an affine neutral mode.
    # The fast controller must remove that mode using only paired neutral
    # soma/traffic observations, without changing the predictor parameters or
    # reading the task instruction during its coefficient fit.
    frozen_before_projection = [value.clone() for value in
                                net.P_traffic + net.P_traffic_bias]
    projected = net.mixed_apical_components(
        instruction, forward["hiddens"], "innovation",
        neutral_projection=True)
    projection = projected["neutral_projection"]
    projected_instruction_error = max(float((left - right).abs().max())
                                      for left, right in zip(
        projected["innovation"], instruction))
    assert projected_instruction_error < 1e-14
    assert projection["post_projection_traffic_rms_ratio"] < 1e-14
    assert projection["max_absolute_post_projection_soma_slope"] < 1e-14
    assert projection["instruction_observations"] == 0
    assert all(torch.equal(before, after) for before, after in zip(
        frozen_before_projection, net.P_traffic + net.P_traffic_bias))

    # Raw and matched controls pay for the identical neutral projection but do
    # not apply its subtractive direction. Raw must remain exactly the mixed
    # apical vector. Matched may borrow only projected innovation's norm.
    projected_raw = net.mixed_apical_components(
        instruction, forward["hiddens"], "raw",
        neutral_projection=True)
    projected_matched = net.mixed_apical_components(
        instruction, forward["hiddens"], "matched",
        neutral_projection=True)
    sham_raw_error = max(float((left - right).abs().max())
                         for left, right in zip(
        projected_raw["used"], projected_raw["raw"]))
    sham_projection_report_error = max(
        abs(projected_raw["neutral_projection"][key]
            - projected_matched["neutral_projection"][key])
        for key in (
            "pre_projection_traffic_rms_ratio",
            "post_projection_traffic_rms_ratio",
            "max_absolute_pre_projection_soma_slope",
            "max_absolute_post_projection_soma_slope",
            "max_positive_post_projection_soma_slope",
            "min_post_projection_soma_slope",
            "max_absolute_correction_slope",
        ))
    sham_matched_norm_errors = []
    sham_matched_direction_errors = []
    for raw, innovation, matched in zip(
            projected_matched["raw"], projected_matched["innovation"],
            projected_matched["matched"]):
        raw_flat = raw.flatten(1)
        innovation_flat = innovation.flatten(1)
        matched_flat = matched.flatten(1)
        sham_matched_norm_errors.append(float((
            (matched_flat.norm(dim=1) - innovation_flat.norm(dim=1)).abs()
            / innovation_flat.norm(dim=1).clamp_min(1e-30)).max()))
        sham_matched_direction_errors.append(float((F.cosine_similarity(
            raw_flat, matched_flat, dim=1) - 1.0).abs().max()))
    assert sham_raw_error == 0.0
    assert sham_projection_report_error == 0.0
    assert max(sham_matched_norm_errors) < 1e-12
    assert max(sham_matched_direction_errors) < 1e-12
    assert projected_raw["neutral_projection"]["instruction_observations"] == 0
    assert all(torch.equal(before, after) for before, after in zip(
        frozen_before_projection, net.P_traffic + net.P_traffic_bias))

    for slope, bias in zip(net.P_traffic, net.P_traffic_bias):
        slope.zero_()
        bias.zero_()
    components = net.mixed_apical_components(
        instruction, forward["hiddens"], "matched")
    norm_errors = []
    direction_errors = []
    for raw, innovation, matched in zip(
            components["raw"], components["innovation"],
            components["matched"]):
        raw_flat = raw.flatten(1)
        innovation_flat = innovation.flatten(1)
        matched_flat = matched.flatten(1)
        norm_errors.append(float((
            (matched_flat.norm(dim=1) - innovation_flat.norm(dim=1)).abs()
            / innovation_flat.norm(dim=1).clamp_min(1e-30)).max()))
        direction_errors.append(float((F.cosine_similarity(
            raw_flat, matched_flat, dim=1) - 1.0).abs().max()))
    assert max(norm_errors) < 1e-12
    assert max(direction_errors) < 1e-12

    # All used signals retain equal independently recomputed local KP products.
    correlation_errors = []
    for rule in ("raw", "matched", "innovation"):
        used = net.mixed_apical_components(
            instruction, forward["hiddens"], rule)["used"]
        forward_directions, _, _, output_weight, _ = (
            net.local_ascent_directions(used, output_error, forward))
        reciprocal, reciprocal_readout = net.reciprocal_feedback_directions(
            used, output_error, forward)
        correlation_errors.extend(float((left - right).abs().max())
                                  for left, right in zip(
            forward_directions[1:], reciprocal[1:]))
        correlation_errors.append(float(
            (reciprocal_readout + output_weight.t()).abs().max()))
    assert max(correlation_errors) < 1e-14

    # Predictor plasticity consumes only the supplied soma/traffic pair: once
    # those are fixed, changing every forward/feedback weight has no effect.
    left = CIFARKPMixedTrafficResNet(**common)
    right = CIFARKPMixedTrafficResNet(**common)
    for left_gain, right_gain, source in zip(
            left.traffic_gain, right.traffic_gain, net.traffic_gain):
        left_gain.copy_(source)
        right_gain.copy_(source)
    fixed_hiddens = [value.detach().clone() for value in forward["hiddens"]]
    for value in right.W + right.Q + [right.W_out, right.R_out]:
        value.add_(torch.randn_like(value))
    left.predictor_step(fixed_hiddens, eta=0.1)
    right.predictor_step(fixed_hiddens, eta=0.1)
    predictor_independence_error = max(float((a - b).abs().max())
                                       for a, b in zip(
        left.P_traffic + left.P_traffic_bias,
        right.P_traffic + right.P_traffic_bias))
    assert predictor_independence_error == 0.0

    closed_left = CIFARKPMixedTrafficResNet(**common)
    closed_right = CIFARKPMixedTrafficResNet(**common)
    for left_gain, right_gain, source in zip(
            closed_left.traffic_gain, closed_right.traffic_gain,
            net.traffic_gain):
        left_gain.copy_(source)
        right_gain.copy_(source)
    for value in (closed_right.W + closed_right.Q
                  + [closed_right.W_out, closed_right.R_out]):
        value.add_(torch.randn_like(value))
    closed_left.predictor_closed_form_fit(fixed_hiddens)
    closed_right.predictor_closed_form_fit(fixed_hiddens)
    closed_form_independence_error = max(float((a - b).abs().max())
                                         for a, b in zip(
        closed_left.P_traffic + closed_left.P_traffic_bias,
        closed_right.P_traffic + closed_right.P_traffic_bias))
    assert closed_form_independence_error == 0.0

    net.traffic_rule = "innovation"
    result = conv_kp_mixed_traffic_step(
        net, x, y, ConvSDILConfig(
            eta=1e-4, eta_output=1e-4, eta_P=0.1, momentum=0.0,
            weight_decay=0.0, learn_A=False, learn_P=True),
        step=0, rule="innovation", predictor_every=16)
    assert math.isfinite(result["loss"]) and result["did_predictor_update"]
    frozen_predictor = [value.clone() for value in
                        net.P_traffic + net.P_traffic_bias]
    frozen_result = conv_kp_mixed_traffic_step(
        net, x, y, ConvSDILConfig(
            eta=1e-4, eta_output=1e-4, eta_P=0.1, momentum=0.0,
            weight_decay=0.0, learn_A=False, learn_P=True),
        step=1, rule="innovation", predictor_every=0)
    assert math.isfinite(frozen_result["loss"])
    assert not frozen_result["did_predictor_update"]
    assert all(torch.equal(before, after) for before, after in zip(
        frozen_predictor, net.P_traffic + net.P_traffic_bias))
    assert all(not value.requires_grad for value in
               net.W + net.Q + net.P_traffic + net.P_traffic_bias
               + [net.W_out, net.R_out, net.b_out])
    projected_result = conv_kp_mixed_traffic_step(
        net, x, y, ConvSDILConfig(
            eta=1e-4, eta_output=1e-4, eta_P=0.1, momentum=0.0,
            weight_decay=0.0, learn_A=False, learn_P=True),
        step=2, rule="innovation", predictor_every=0,
        neutral_projection=True)
    assert math.isfinite(projected_result["loss"])
    assert projected_result["neutral_projection"] is not None
    assert net.mixed_elementwise_ops_per_example("matched") > (
        net.mixed_elementwise_ops_per_example("raw"))
    return {
        "kp_traffic_zero_limit_error": max(zero_errors),
        "kp_traffic_ratio_error": ratio_error,
        "kp_traffic_exact_predictor_error": exact_predictor_error,
        "kp_traffic_closed_form_predictor_error": closed_form_error,
        "kp_traffic_closed_form_residual_ratio": closed_form[
            "residual_traffic_rms_ratio"],
        "kp_traffic_closed_form_residual_slope": closed_form[
            "max_absolute_residual_soma_slope"],
        "kp_traffic_projected_instruction_error": projected_instruction_error,
        "kp_traffic_projected_residual_ratio": projection[
            "post_projection_traffic_rms_ratio"],
        "kp_traffic_projected_residual_slope": projection[
            "max_absolute_post_projection_soma_slope"],
        "kp_traffic_sham_raw_error": sham_raw_error,
        "kp_traffic_sham_projection_report_error": (
            sham_projection_report_error),
        "kp_traffic_sham_matched_norm_error": max(
            sham_matched_norm_errors),
        "kp_traffic_sham_matched_direction_error": max(
            sham_matched_direction_errors),
        "kp_traffic_matched_norm_error": max(norm_errors),
        "kp_traffic_matched_direction_error": max(direction_errors),
        "kp_traffic_reciprocal_correlation_error": max(correlation_errors),
        "kp_traffic_predictor_parameter_independence_error": (
            predictor_independence_error),
        "kp_traffic_closed_form_parameter_independence_error": (
            closed_form_independence_error),
    }


def apical_learning_checks():
    torch.manual_seed(11)
    net = CIFARSDILResNet(depth=8, base_width=2, seed=6)
    x = torch.randn(8, 3, 32, 32)
    y = torch.arange(8) % 10
    clean = net.forward(x, return_cache=True)
    output_signal = (torch.softmax(clean["logits"], dim=1)
                     - F.one_hot(y, 10))
    prediction, _, _ = net.apical_components(
        output_signal, clean["hiddens"], use_residual=True)
    targets = [torch.randn_like(value) * 0.01 for value in prediction]
    before = sum(float((target - value).square().sum())
                 for target, value in zip(targets, prediction))
    net.calibrate_apical(
        output_signal, clean["hiddens"], prediction, targets, eta=0.1)
    after_prediction, _, _ = net.apical_components(
        output_signal, clean["hiddens"], use_residual=True)
    after = sum(float((target - value).square().sum())
                for target, value in zip(targets, after_prediction))
    assert after < before

    gated = CIFARSDILResNet(
        depth=8, base_width=2, seed=6, vectorizer_mode="channel_gated")
    gated_clean = gated.forward(x)
    gated_signal = (torch.softmax(gated_clean["logits"], dim=1)
                    - F.one_hot(y, 10))
    gated_prediction, _, _ = gated.apical_components(
        gated_signal, gated_clean["hiddens"], use_residual=True)
    gated_targets = [torch.randn_like(value) * 0.01 for value in gated_prediction]
    gated_before = sum(float((target - value).square().sum())
                       for target, value in zip(gated_targets, gated_prediction))
    gated.calibrate_apical(
        gated_signal, gated_clean["hiddens"], gated_prediction,
        gated_targets, eta=0.1)
    gated_after_prediction, _, _ = gated.apical_components(
        gated_signal, gated_clean["hiddens"], use_residual=True)
    gated_after = sum(float((target - value).square().sum())
                      for target, value in zip(gated_targets, gated_after_prediction))
    assert gated_after < gated_before
    shifted_hidden = [torch.roll(value, shifts=(3, -2), dims=(2, 3))
                      for value in gated_clean["hiddens"]]
    shifted_instruction, _, _ = gated.apical_components(
        gated_signal, shifted_hidden, use_residual=True)
    original_instruction, _, _ = gated.apical_components(
        gated_signal, gated_clean["hiddens"], use_residual=True)
    assert all(torch.allclose(
        shifted, torch.roll(original, shifts=(3, -2), dims=(2, 3)))
        for shifted, original in zip(shifted_instruction, original_instruction))
    spatial_56 = CIFARSDILResNet(depth=56, vectorizer_mode="spatial_template")
    gated_56 = CIFARSDILResNet(depth=56, vectorizer_mode="channel_gated")
    assert spatial_56.n_vectorizer_parameters == 5_324_800
    assert gated_56.n_vectorizer_parameters == 40_640

    predictor_net = CIFARSDILResNet(depth=8, base_width=2, seed=8)
    hiddens = [torch.randn(64, *shape) for shape in predictor_net.hidden_shapes]
    initial = predictor_net.predictor_step(hiddens, eta=0.1, nuisance_scale=0.5)
    final = initial
    for _ in range(60):
        final = predictor_net.predictor_step(hiddens, eta=0.1, nuisance_scale=0.5)
    assert final < initial * 1e-3

    weights_before = [weight.clone() for weight in net.W]
    result = conv_local_step(
        net, x[:2], y[:2],
        ConvSDILConfig(
            eta=1e-3, eta_A=1e-3, momentum=0.0, weight_decay=0.0,
            pert_every=1),
        step=0, generator=torch.Generator(device="cpu").manual_seed(7))
    assert result["did_perturb"] and result["calibration"] is not None
    assert torch.isfinite(torch.tensor(list(
        value for key, value in result.items()
        if isinstance(value, float) and key != "predictor_mse"))).all()
    assert any(not torch.equal(before_weight, after_weight)
               for before_weight, after_weight in zip(weights_before, net.W))
    assert all(not parameter.requires_grad
               for parameter in net.W + [net.W_out, net.b_out])

    gated_weights_before = [value.clone() for value in gated.A + gated.A_gate]
    gated_result = conv_local_step(
        gated, x[:2], y[:2],
        ConvSDILConfig(
            eta=1e-3, eta_A=1e-3, momentum=0.0, weight_decay=0.0,
            pert_every=1, apical_calibration_mode="channel_subspace"),
        step=0, generator=torch.Generator(device="cpu").manual_seed(17))
    assert gated_result["did_perturb"]
    assert all(torch.isfinite(torch.tensor(value)) for value in
               gated_result["calibration"].values())
    assert any(not torch.equal(before, after) for before, after in zip(
        gated_weights_before, gated.A + gated.A_gate))
    return {"apical_mse_ratio": after / before,
            "gated_apical_mse_ratio": gated_after / gated_before,
            "gated_vectorizer_parameter_reduction": (
                spatial_56.n_vectorizer_parameters
                / gated_56.n_vectorizer_parameters),
            "predictor_mse_ratio": final / initial}


def main():
    architecture_checks()
    perturbation_checks()
    report = exact_local_gradient_check()
    report.update(exact_batchnorm_local_gradient_check())
    report.update(perturbation_estimator_check())
    report.update(channel_subspace_estimator_check())
    report.update(vectorizer_subspace_estimator_check())
    report.update(hierarchical_feedback_checks())
    report.update(hierarchical_parameter_calibration_checks())
    report.update(normalized_response_mirror_checks())
    report.update(kolen_pollack_checks())
    report.update(kp_mixed_traffic_checks())
    report.update(apical_learning_checks())
    print(report)
    print("ALL CONVOLUTIONAL LOCAL-ELIGIBILITY CHECKS PASSED")


if __name__ == "__main__":
    main()