Listing: Listing

import math
import torch
from torch.utils.data import DataLoader

@torch.no_grad()
def forward_explicit_noise(model, x, cue, sigma, generator):
    if sigma < 0:
        raise ValueError("sigma must be nonnegative")

    h = torch.cat([x, cue], dim=1)
    for layer in (model.fc1, model.fc2, model.fc3):
        h = torch.relu(layer(h))
        if sigma > 0:
            noise = torch.randn(
                h.shape,
                dtype=h.dtype,
                device=h.device,
                generator=generator,
            )
            h = h + sigma * noise
    return model.head(h)

@torch.no_grad()
def paired_cue_probe(model, x, y, seed=10001, batch_size=1000):
    if len(x) == 0:
        raise ValueError("Empty test set")

    was_training = model.training
    model.eval()
    generator = torch.Generator(device=x.device)
    generator.manual_seed(int(seed))
    permutation = torch.tensor(
        [2, 3, 0, 1], device=y.device, dtype=torch.long
    )
    counts = {
        "correct_visual": 0,
        "wrong_visual": 0,
        "correct_voice": 0,
        "zero_visual": 0,
        "midpoint_visual": 0,
        "midpoint_voice": 0,
        "random_visual": 0,
    }

    try:
        for start in range(0, len(x), batch_size):
            xb = x[start:start + batch_size]
            yb = y[start:start + batch_size]
            y_voice = permutation[yb]
            n = len(xb)

            cv = xb.new_tensor([1.0, 0.0]).expand(n, 2)
            ca = xb.new_tensor([0.0, 1.0]).expand(n, 2)
            cz = xb.new_zeros((n, 2))
            cm = xb.new_full((n, 2), 0.5)

            # Independent, normalized per-example random cue.
            alpha = torch.rand(
                (n, 1),
                dtype=xb.dtype,
                device=xb.device,
                generator=generator,
            )
            cr = torch.cat([alpha, 1.0 - alpha], dim=1)

            def predict(cue):
                return forward_explicit_noise(
                    model, xb, cue, 0.0, generator
                ).argmax(dim=1)

            pv = predict(cv)
            pa = predict(ca)
            pz = predict(cz)
            pm = predict(cm)
            pr = predict(cr)

            counts["correct_visual"] += int((pv == yb).sum())
            counts["wrong_visual"] += int((pa == yb).sum())
            counts["correct_voice"] += int(
                (pa == y_voice).sum()
            )
            counts["zero_visual"] += int((pz == yb).sum())
            counts["midpoint_visual"] += int((pm == yb).sum())
            counts["midpoint_voice"] += int(
                (pm == y_voice).sum()
            )
            counts["random_visual"] += int((pr == yb).sum())
    finally:
        model.train(was_training)

    result = {key: value / len(x) for key, value in counts.items()}
    result["cue_drop"] = (
        result["correct_visual"] - result["wrong_visual"]
    )
    result["zero_cue_drop"] = (
        result["correct_visual"] - result["zero_visual"]
    )
    result["balanced_midpoint"] = 0.5 * (
        result["midpoint_visual"] + result["midpoint_voice"]
    )
    assert result["balanced_midpoint"] <= 0.5 + 1e-12
    return result

@torch.no_grad()
def stress_mc(
    model, x, y, sigma, seeds,
    batch_size=1000, delta=0.05
):
    seeds = list(seeds)
    if len(x) == 0 or len(seeds) == 0:
        raise ValueError("Nonempty inputs and seeds are required")
    if not 0.0 < delta < 1.0:
        raise ValueError("delta must lie in (0, 1)")

    was_training = model.training
    model.eval()
    replicate_means = []
    try:
        for seed in seeds:
            generator = torch.Generator(device=x.device)
            generator.manual_seed(int(seed))
            correct = 0
            for start in range(0, len(x), batch_size):
                xb = x[start:start + batch_size]
                yb = y[start:start + batch_size]
                cue = xb.new_tensor([1.0, 0.0]).expand(len(xb), 2)
                pred = forward_explicit_noise(
                    model, xb, cue, sigma, generator
                ).argmax(dim=1)
                correct += int((pred == yb).sum())
            replicate_means.append(correct / len(x))
    finally:
        model.train(was_training)

    n_noise_outcomes = len(x) * len(seeds)
    radius = math.sqrt(
        math.log(2.0 / delta) / (2.0 * n_noise_outcomes)
    )
    return {
        "mean": sum(replicate_means) / len(replicate_means),
        "replicate_means": replicate_means,
        # Conditional on this fixed held-out input set.
        "conditional_mc_hoeffding_radius": radius,
    }