Listing: Listing

import copy
import math
from fractions import Fraction
import torch
from torch import nn

def fraction01(value):
    value = float(value)
    if not math.isfinite(value) or not 0.0 <= value <= 1.0:
        raise ValueError("fraction must be finite and in [0, 1]")
    return Fraction(str(value))

def ceil_fraction(value):
    return (
        value.numerator + value.denominator - 1
    ) // value.denominator

def generator(seed):
    return torch.Generator(device="cpu").manual_seed(int(seed))

class CheckedNetwork(nn.Module):
    def __init__(self, widths=(512, 512, 256), seed=0):
        super().__init__()
        self.widths = tuple(widths)
        # Constructor does not advance the caller's global stream.
        with torch.random.fork_rng(devices=[]):
            torch.manual_seed(int(seed))
            self.fc1 = nn.Linear(2, widths[0])
            self.fc2 = nn.Linear(widths[0], widths[1])
            self.fc3 = nn.Linear(widths[1], widths[2])
            self.fc4 = nn.Linear(widths[2], 4)
        self.a = 1.0
        self.activation = "relu"

    def forward(self, x, sigma=0.0, noise_generator=None):
        if not 0.0 <= self.a <= 1.0:
            raise ValueError("damping outside [0, 1]")
        if sigma < 0 or not math.isfinite(float(sigma)):
            raise ValueError("invalid stress")
        if self.activation not in ("relu", "tanh"):
            raise ValueError("unsupported activation")
        phi = (
            torch.relu if self.activation == "relu"
            else torch.tanh
        )
        h = x
        for layer in (self.fc1, self.fc2, self.fc3):
            h = phi(layer(h))
            if sigma > 0:
                if noise_generator is None:
                    raise ValueError("explicit noise generator required")
                h = h + sigma * torch.randn(
                    h.shape, dtype=h.dtype,
                    generator=noise_generator
                )
            h = self.a * h
        return self.fc4(h)

class MaskLedger:
    def __init__(self, model, masks=None):
        self.model = model
        weights = {
            name: p for name, p in model.named_parameters()
            if name.endswith("weight")
        }
        if masks is not None and set(masks) != set(weights):
            raise ValueError("mask keys do not match weights")
        self.masks = {}
        for name, p in weights.items():
            if masks is None:
                mask = torch.ones_like(p, dtype=torch.bool)
            else:
                raw = masks[name]
                if raw.shape != p.shape:
                    raise ValueError("mask shape mismatch")
                if not bool(((raw == 0) | (raw == 1)).all()):
                    raise ValueError("nonbinary mask")
                mask = raw.bool().clone()
            self.masks[name] = mask
        self.apply()

    @torch.no_grad()
    def apply(self):
        for name, mask in self.masks.items():
            p = self.model.get_parameter(name)
            p.mul_(mask.to(p.dtype))

    def density(self):
        kept = sum(int(m.sum()) for m in self.masks.values())
        total = sum(m.numel() for m in self.masks.values())
        return kept / total

    @torch.no_grad()
    def keep_counts(self, counts):
        # Validate every requested deletion before changing state.
        if set(counts) != set(self.masks):
            raise ValueError("incomplete count specification")
        for name, k in counts.items():
            p = self.model.get_parameter(name)
            active = int(self.masks[name].sum())
            if not isinstance(k, int) or not 0 <= k <= active:
                raise ValueError("invalid count or implicit regrowth")
            if not bool(torch.isfinite(p).all()):
                raise ValueError("non-finite weights")

        for name, k in counts.items():
            p = self.model.get_parameter(name)
            old = self.masks[name].reshape(-1)
            eligible = torch.nonzero(old).flatten()
            values = p.detach().reshape(-1)[eligible].abs()
            order = torch.argsort(
                values, descending=True, stable=True
            )
            new = torch.zeros_like(old)
            new[eligible[order[:k]]] = True
            self.masks[name] = new.reshape_as(p)
            assert int(new.sum()) == k
            assert not bool((new & ~old).any())
        self.apply()

    def prune_to_density(self, density):
        d = fraction01(density)
        counts = {
            name: ceil_fraction(d * mask.numel())
            for name, mask in self.masks.items()
        }
        self.keep_counts(counts)

    def secondary_prune(self, fraction):
        f = fraction01(fraction)
        counts = {}
        for name, mask in self.masks.items():
            active = int(mask.sum())
            amount = f * active
            remove = amount.numerator // amount.denominator
            counts[name] = active - remove
        self.keep_counts(counts)

    @torch.no_grad()
    def regrow(self, fraction, scores, init_scale, seed):
        r = fraction01(fraction)
        if init_scale < 0 or not math.isfinite(init_scale):
            raise ValueError("invalid initialization scale")
        plan = {}
        for name, mask in self.masks.items():
            eligible = torch.nonzero(~mask.reshape(-1)).flatten()
            amount = r * len(eligible)
            k = amount.numerator // amount.denominator
            plan[name] = (eligible, k)

        # Zero regrowth changes neither parameters nor random streams.
        if not any(k for _, k in plan.values()):
            return {name: 0 for name in plan}

        for name, (eligible, k) in plan.items():
            score = scores[name]
            if score.shape != self.masks[name].shape:
                raise ValueError("score shape mismatch")
            if not bool(torch.isfinite(score).all()):
                raise ValueError("non-finite score")

        noise = generator(seed)
        result = {}
        for name, (eligible, k) in plan.items():
            p = self.model.get_parameter(name)
            values = scores[name].reshape(-1)[eligible]
            order = torch.argsort(
                values, descending=True, stable=True
            )
            selected = eligible[order[:k]]
            self.masks[name].view(-1)[selected] = True
            p.view(-1)[selected] = init_scale * torch.randn(
                (k,), dtype=p.dtype, generator=noise
            )
            result[name] = k
        self.apply()
        return result

def gradient_scores(model, ledger, x, y, batch_size=128,
                    num_batches=30, order_seed=1000):
    if len(x) == 0 or len(x) != len(y):
        raise ValueError("invalid training sample")
    if batch_size <= 0 or num_batches <= 0:
        raise ValueError("positive batch controls required")
    scores = {
        name: torch.zeros_like(model.get_parameter(name))
        for name in ledger.masks
    }
    was_training = model.training
    order = torch.randperm(len(x), generator=generator(order_seed))
    used = 0
    try:
        ledger.apply()
        model.train()
        for start in range(0, len(x), batch_size):
            if used >= num_batches:
                break
            idx = order[start:start + batch_size]
            model.zero_grad(set_to_none=True)
            # Stored inactive weights are zero; no mask multiplication
            # inside autograd blocks their ambient derivatives.
            logits = model(x[idx], sigma=0.0)
            nn.functional.cross_entropy(logits, y[idx]).backward()
            for name, mask in ledger.masks.items():
                grad = model.get_parameter(name).grad
                if grad is None:
                    raise ValueError("missing ambient gradient")
                scores[name].add_(grad.detach().abs() * (~mask))
            used += 1
    finally:
        model.zero_grad(set_to_none=True)
        model.train(was_training)
    if used == 0:
        raise ValueError("no gradient batches")
    return {name: value / used for name, value in scores.items()}

def fit(model, ledger, x, y, epochs, lr,
        initial_sigma=0.0, final_sigma=0.0,
        order_seed=2000, noise_seed=3000, batch_size=128):
    if not isinstance(epochs, int) or epochs < 0:
        raise ValueError("invalid epoch count")
    if len(x) == 0 or len(x) != len(y) or batch_size <= 0:
        raise ValueError("invalid training data")
    if lr <= 0 or not math.isfinite(lr):
        raise ValueError("invalid learning rate")
    if initial_sigma < 0 or final_sigma < 0:
        raise ValueError("invalid training-noise schedule")
    if epochs == 1 and initial_sigma != final_sigma:
        raise ValueError("one epoch cannot realize two endpoints")

    order_rng, noise_rng = generator(order_seed), generator(noise_seed)
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    was_training = model.training
    losses, steps, examples = [], 0, 0
    try:
        for epoch in range(epochs):
            progress = epoch / (epochs - 1) if epochs > 1 else 0.0
            sigma = initial_sigma + progress * (
                final_sigma - initial_sigma
            )
            order = torch.randperm(len(x), generator=order_rng)
            model.train()
            total_loss = 0.0
            for start in range(0, len(x), batch_size):
                idx = order[start:start + batch_size]
                optimizer.zero_grad(set_to_none=True)
                logits = model(x[idx], sigma, noise_rng)
                loss = nn.functional.cross_entropy(logits, y[idx])
                loss.backward()
                optimizer.step()
                ledger.apply()
                total_loss += float(loss.detach()) * len(idx)
                steps += 1
                examples += len(idx)
            losses.append(total_loss / len(x))
    finally:
        model.zero_grad(set_to_none=True)
        model.train(was_training)

    return {
        "losses": losses,
        "optimizer_steps": steps,
        "examples_processed": examples,
        "optimizer_state": copy.deepcopy(optimizer.state_dict()),
        "order_rng_state": order_rng.get_state().clone(),
        "noise_rng_state": noise_rng.get_state().clone(),
        "schedule": (initial_sigma, final_sigma, epochs),
        "optimizer_reset": True,
    }

def phase_snapshot(model, ledger):
    # Valid at a boundary where a later fit creates a fresh optimizer.
    return {
        "weights_biases": {
            key: value.detach().clone()
            for key, value in model.state_dict().items()
        },
        "masks": {
            key: value.clone() for key, value in ledger.masks.items()
        },
        "a": model.a,
        "activation": model.activation,
        "widths": model.widths,
        "training": model.training,
    }

def clone_phase(snapshot):
    model = CheckedNetwork(snapshot["widths"], seed=0)
    model.load_state_dict(snapshot["weights_biases"])
    model.a = snapshot["a"]
    model.activation = snapshot["activation"]
    model.train(snapshot["training"])
    ledger = MaskLedger(model, snapshot["masks"])
    return model, ledger

@torch.no_grad()
def budget(pre, model, ledger):
    state = model.state_dict()
    dw1 = dw2 = db1 = 0.0
    nw = nb = flips = 0
    for name, old in pre["weights_biases"].items():
        delta = (state[name] - old).double()
        if name.endswith("weight"):
            dw1 += float(delta.abs().sum())
            dw2 += float(delta.square().sum())
            nw += delta.numel()
            flips += int(
                (ledger.masks[name] != pre["masks"][name]).sum()
            )
        elif name.endswith("bias"):
            db1 += float(delta.abs().sum())
            nb += delta.numel()
    return {
        "d_w": dw1 / nw,
        "d_2_code": math.sqrt(dw2) / nw,
        "d_2_rms": math.sqrt(dw2 / nw),
        "d_b": db1 / nb,
        "d_mask": flips / nw,
        "d_a": abs(model.a - pre["a"]),
        "activation_changed": (
            model.activation != pre["activation"]
        ),
    }

@torch.no_grad()
def assess(model, x, y, input_noise, stress, seeds,
           batch_size=1000):
    if len(x) == 0 or len(x) != len(y) or not seeds:
        raise ValueError("nonempty sample and seed bank required")
    if input_noise < 0 or stress < 0 or batch_size <= 0:
        raise ValueError("invalid evaluation controls")
    was_training = model.training
    all_correct = []
    try:
        model.eval()
        for seed in seeds:
            input_rng = generator(2 * int(seed))
            hidden_rng = generator(2 * int(seed) + 1)
            correct = []
            for start in range(0, len(x), batch_size):
                xb = x[start:start + batch_size]
                yb = y[start:start + batch_size]
                if input_noise > 0:
                    xb = xb + input_noise * torch.randn(
                        xb.shape, dtype=xb.dtype, generator=input_rng
                    )
                pred = model(xb, stress, hidden_rng).argmax(dim=1)
                correct.append((pred == yb).double())
            all_correct.append(torch.cat(correct))
    finally:
        model.train(was_training)
    values = torch.stack(all_correct)
    return {
        "accuracy": float(values.mean()),
        "per_example_mean": values.mean(dim=0),
        "correctness": values,
        "seeds": tuple(seeds),
    }

def evaluate_acute_and_challenge(
    model, ledger, pre, x, y, seeds,
    input_noise=1.0, stress=0.5, challenge_fraction=0.4
):
    ledger.apply()
    displacement = budget(pre, model, ledger)
    acute_density = ledger.density()
    acute = assess(
        model, x, y, input_noise, stress, seeds
    )

    challenged_model, challenged_ledger = clone_phase(
        phase_snapshot(model, ledger)
    )
    challenged_ledger.secondary_prune(challenge_fraction)
    challenged = assess(
        challenged_model, x, y, input_noise, stress, seeds
    )

    # The acute checkpoint is not mutated by the challenge.
    return {
        "budget_acute": displacement,
        "density_acute": acute_density,
        "density_challenged": challenged_ledger.density(),
        "accuracy_acute": acute["accuracy"],
        "accuracy_challenged": challenged["accuracy"],
        "challenge_drop": acute["accuracy"] - challenged["accuracy"],
        "paired_per_example_drop": (
            acute["per_example_mean"]
            - challenged["per_example_mean"]
        ),
    }