Listing: Listing

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

def finite(value, name):
    value = float(value)
    if not math.isfinite(value):
        raise ValueError(f"{name} must be finite")
    return value

def fraction01(value):
    value = finite(value, "fraction")
    if not 0.0 <= value <= 1.0:
        raise ValueError("fraction outside [0, 1]")
    return Fraction(str(value))

def floor_count(value, n):
    q = fraction01(value) * int(n)
    return q.numerator // q.denominator

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

class LongitudinalNetwork(nn.Module):
    def __init__(self, widths=(512, 512, 256), seed=0):
        super().__init__()
        if len(widths) != 3 or any(int(w) <= 0 for w in widths):
            raise ValueError("three positive widths required")
        self.widths = tuple(int(w) for w in widths)
        with torch.random.fork_rng(devices=[]):
            torch.manual_seed(int(seed))
            self.fc1 = nn.Linear(2, self.widths[0])
            self.fc2 = nn.Linear(self.widths[0], self.widths[1])
            self.fc3 = nn.Linear(self.widths[1], self.widths[2])
            self.fc4 = nn.Linear(self.widths[2], 4)
        self.a, self.g, self.p = 1.0, 1.0, 0.0
        self.decay_rate = 0.0
        self.activation = "relu"

    def effective_gain(self):
        a, g, p = (finite(self.a, "a"), finite(self.g, "g"),
                   finite(self.p, "p"))
        if not 0 <= a <= 1 or g < 0 or not 0 <= p <= 1:
            raise ValueError("invalid modulation")
        return g - p * max(g - 1.05, 0.0)

    def reset_transient_controls(self, remove_protection=False):
        self.a, self.g, self.activation = 1.0, 1.0, "relu"
        if remove_protection:
            self.p, self.decay_rate = 0.0, 0.0

    def decay_protection(self):
        d = finite(self.decay_rate, "decay_rate")
        if d < 0:
            raise ValueError("negative protection decrement")
        self.p = max(0.0, self.p - d)

    def forward(self, x, beta=0.0, sigma=0.0,
                noise_generator=None, return_hidden=False):
        beta, sigma = finite(beta, "beta"), finite(sigma, "sigma")
        if sigma < 0:
            raise ValueError("negative stress")
        if self.activation not in ("relu", "tanh"):
            raise ValueError("unsupported activation")
        phi = torch.relu if self.activation == "relu" else torch.tanh
        q = self.a * self.effective_gain()
        h, hidden = 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 stream required")
                h = h + sigma * torch.randn(
                    h.shape, dtype=h.dtype, generator=noise_generator
                )
            h = q * (h + beta * (1.0 - 0.3 * self.p))
            h = h - 0.15 * self.p
            hidden.append(h)
        z = self.fc4(h)
        return (z, hidden) if return_hidden else z

class ScarLedger:
    def __init__(self, model, masks=None, scars=None):
        self.model = model
        weights = {
            n: p for n, p in model.named_parameters()
            if n.endswith("weight")
        }
        for supplied in (masks, scars):
            if supplied is not None and set(supplied) != set(weights):
                raise ValueError("incomplete structural state")
        self.masks, self.scars = {}, {}
        for name, p in weights.items():
            m = torch.ones_like(p, dtype=torch.bool) if masks is None \
                else self._binary(masks[name], p.shape)
            s = torch.zeros_like(p, dtype=torch.bool) if scars is None \
                else self._binary(scars[name], p.shape)
            if bool((m & s).any()):
                raise ValueError("active coordinate intersects scar")
            self.masks[name], self.scars[name] = m, s
        self.apply()

    @staticmethod
    def _binary(value, shape):
        if value.shape != shape:
            raise ValueError("structural shape mismatch")
        if not bool(((value == 0) | (value == 1)).all()):
            raise ValueError("nonbinary structural state")
        return value.bool().clone().contiguous()

    def validate(self):
        for name, mask in self.masks.items():
            p = self.model.get_parameter(name)
            if mask.shape != p.shape or self.scars[name].shape != p.shape:
                raise ValueError("structural shape mismatch")
            if bool((mask & self.scars[name]).any()):
                raise ValueError("structural invariant violated")
            if not bool(torch.isfinite(p).all()):
                raise ValueError("non-finite weights")

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

    def counts(self):
        total = sum(m.numel() for m in self.masks.values())
        active = sum(int(m.sum()) for m in self.masks.values())
        scarred = sum(int(s.sum()) for s in self.scars.values())
        return {
            "active": active, "scarred": scarred, "total": total,
            "active_density": active / total,
            "scar_density": scarred / total,
            "repair_eligible": total - active - scarred,
        }

    @torch.no_grad()
    def delete_fraction(self, fraction, permanent=False,
                        minimum_one=False):
        fraction01(fraction)
        self.validate()
        plans = {}
        for name, mask in self.masks.items():
            eligible = torch.nonzero(
                (mask & ~self.scars[name]).reshape(-1)
            ).flatten()
            n = len(eligible)
            k = floor_count(fraction, n)
            if minimum_one and n:
                k = min(n, max(1, k))
            values = self.model.get_parameter(name).detach() \
                .reshape(-1)[eligible].abs()
            order = torch.argsort(values, stable=True)
            plans[name] = eligible[order[:k]]
        for name, selected in plans.items():
            if permanent:
                self.scars[name].view(-1)[selected] = True
            self.masks[name].view(-1)[selected] = False
        self.apply()
        return {n: len(v) for n, v in plans.items()}

    @torch.no_grad()
    def initial_scars(self, requested_fraction, placement_seed):
        # Counts use the original layer size, like the source.
        # All unscarred positions are eligible, including inactive ones.
        fraction01(requested_fraction)
        self.validate()
        noise, plans = rng(placement_seed), {}
        for name, mask in self.masks.items():
            eligible = torch.nonzero(
                (~self.scars[name]).reshape(-1)
            ).flatten()
            k = floor_count(requested_fraction, mask.numel())
            if k > len(eligible):
                raise ValueError("requested scars exceed eligibility")
            order = torch.randperm(len(eligible), generator=noise)
            plans[name] = eligible[order[:k]]
        for name, selected in plans.items():
            self.scars[name].view(-1)[selected] = True
            self.masks[name].view(-1)[selected] = False
        self.apply()
        return {n: len(v) for n, v in plans.items()}

    @torch.no_grad()
    def regrow(self, fraction, scores, init_scale, seed):
        fraction01(fraction)
        init_scale = finite(init_scale, "init_scale")
        if init_scale < 0:
            raise ValueError("negative initialization scale")
        self.validate()
        plans = {}
        for name, mask in self.masks.items():
            eligible = torch.nonzero(
                (~mask & ~self.scars[name]).reshape(-1)
            ).flatten()
            k = floor_count(fraction, len(eligible))
            if k == 0:
                plans[name] = eligible[:0]
                continue
            score = scores[name]
            if score.shape != mask.shape or not bool(
                torch.isfinite(score).all()
            ):
                raise ValueError("invalid repair scores")
            order = torch.argsort(
                score.reshape(-1)[eligible],
                descending=True, stable=True
            )
            plans[name] = eligible[order[:k]]
        if not any(len(v) for v in plans.values()):
            return {n: 0 for n in plans}
        noise = rng(seed)
        for name, selected in plans.items():
            p = self.model.get_parameter(name)
            self.masks[name].view(-1)[selected] = True
            p.view(-1)[selected] = init_scale * torch.randn(
                (len(selected),), dtype=p.dtype, generator=noise
            )
        self.apply()
        return {n: len(v) for n, v in plans.items()}

def snapshot(model, ledger, phase, streams=None):
    return {
        "phase": str(phase),
        "state_dict": {
            k: v.detach().clone() for k, v in model.state_dict().items()
        },
        "masks": {k: v.clone() for k, v in ledger.masks.items()},
        "scars": {k: v.clone() for k, v in ledger.scars.items()},
        "controls": {
            "a": model.a, "g": model.g, "p": model.p,
            "decay_rate": model.decay_rate,
            "activation": model.activation,
        },
        "widths": model.widths,
        "training": model.training,
        "streams": {
            k: v.get_state().clone() for k, v in (streams or {}).items()
        },
        "optimizer_policy": "fresh_optimizer_at_next_fit",
    }

def restore(saved):
    model = LongitudinalNetwork(saved["widths"], seed=0)
    model.load_state_dict(saved["state_dict"])
    for name, value in saved["controls"].items():
        setattr(model, name, value)
    model.train(saved["training"])
    ledger = ScarLedger(model, saved["masks"], saved["scars"])
    streams = {}
    for name, state in saved["streams"].items():
        streams[name] = rng(0)
        streams[name].set_state(state.clone())
    return model, ledger, streams

def fit(model, ledger, x, y, epochs, lr, order_seed, noise_seed,
        sigma_start=0.0, sigma_end=0.0, batch_size=128):
    if not isinstance(epochs, int) or epochs < 0:
        raise ValueError("invalid epoch count")
    lr = finite(lr, "learning rate")
    s0, s1 = finite(sigma_start, "sigma_start"), \
             finite(sigma_end, "sigma_end")
    if lr <= 0 or min(s0, s1) < 0:
        raise ValueError("invalid training controls")
    if len(x) == 0 or len(x) != len(y) or batch_size <= 0:
        raise ValueError("invalid training data")
    if epochs == 1 and s0 != s1:
        raise ValueError("one epoch cannot realize two endpoints")
    order, noise = rng(order_seed), rng(noise_seed)
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    was_training, steps, examples = model.training, 0, 0
    try:
        for epoch in range(epochs):
            u = epoch / (epochs - 1) if epochs > 1 else 0.0
            sigma = s0 + u * (s1 - s0)
            permutation = torch.randperm(len(x), generator=order)
            model.train()
            for start in range(0, len(x), batch_size):
                idx = permutation[start:start + batch_size]
                optimizer.zero_grad(set_to_none=True)
                loss = nn.functional.cross_entropy(
                    model(x[idx], sigma=sigma, noise_generator=noise),
                    y[idx]
                )
                loss.backward()
                optimizer.step()
                ledger.apply()
                steps += 1
                examples += len(idx)
    finally:
        model.zero_grad(set_to_none=True)
        model.train(was_training)
    return {
        "optimizer_steps": steps, "examples_processed": examples,
        "optimizer_reset": True,
        "order_state": order.get_state().clone(),
        "noise_state": noise.get_state().clone(),
        "optimizer_state": copy.deepcopy(optimizer.state_dict()),
    }

def ambient_scores(model, ledger, x, y, order_seed,
                   num_batches=30, batch_size=128):
    if len(x) == 0 or len(x) != len(y):
        raise ValueError("invalid scoring data")
    if num_batches <= 0 or batch_size <= 0:
        raise ValueError("positive scoring controls required")
    scores = {
        n: torch.zeros_like(model.get_parameter(n))
        for n in ledger.masks
    }
    incoming_gain, was_training = model.g, model.training
    permutation = torch.randperm(len(x), generator=rng(order_seed))
    used = 0
    try:
        ledger.apply()
        model.g = 1.0
        model.train()
        for start in range(0, len(x), batch_size):
            if used >= num_batches:
                break
            idx = permutation[start:start + batch_size]
            model.zero_grad(set_to_none=True)
            loss = nn.functional.cross_entropy(
                model(x[idx], beta=0.0, sigma=0.0), y[idx]
            )
            loss.backward()
            for name in scores:
                grad = model.get_parameter(name).grad
                if grad is None:
                    raise ValueError("missing ambient derivative")
                eligible = ~ledger.masks[name] & ~ledger.scars[name]
                scores[name].add_(grad.detach().abs() * eligible)
            used += 1
    finally:
        model.zero_grad(set_to_none=True)
        model.g = incoming_gain
        model.train(was_training)
    return {n: v / used for n, v in scores.items()}

@torch.no_grad()
def assess(model, x, y, beta, sigma, noise_seed,
           activation_clip, batch_size=1000):
    if len(x) == 0 or len(x) != len(y) or batch_size <= 0:
        raise ValueError("invalid assessment block")
    cap = finite(activation_clip, "activation_clip")
    if cap <= 0:
        raise ValueError("positive diagnostic cap required")
    noise, was_training = rng(noise_seed), model.training
    correct, diagnostic = [], []
    try:
        model.eval()
        for start in range(0, len(x), batch_size):
            xb = x[start:start + batch_size]
            yb = y[start:start + batch_size]
            z = model(xb, beta, sigma, noise)
            _, clean_hidden = model(
                xb, beta=0.0, sigma=0.0, return_hidden=True
            )
            h = sum(
                v.abs().mean(dim=1) for v in clean_hidden
            ) / len(clean_hidden)
            correct.append((z.argmax(dim=1) == yb).double())
            diagnostic.append(h.clamp(max=cap).double())
    finally:
        model.train(was_training)
    c, h = torch.cat(correct), torch.cat(diagnostic)
    return {
        "q_hat": float(c.mean()),
        "h_hat_clipped": float(h.mean()),
        "correctness": c,
        "n": len(x),
        "beta": float(beta), "sigma": float(sigma),
        "noise_seed": int(noise_seed),
        "activation_clip": cap,
        "diagnostic": "protected_clean_equal_layer_mean",
    }

def radius(n, delta, interval_length=1.0):
    delta = finite(delta, "delta")
    if n < 1 or not 0 < delta < 1:
        raise ValueError("invalid confidence controls")
    return interval_length * math.sqrt(
        math.log(2.0 / delta) / (2.0 * n)
    )

def controller_cycle(model, ledger, x_fresh, y_fresh, beta, sigma,
                     noise_seed, q_star, delta_q, delta_h,
                     activation_clip, base_fraction=0.05,
                     minimum_one=False):
    q_star = finite(q_star, "q_star")
    if not 0 < q_star < 1:
        raise ValueError("invalid threshold")
    fraction01(base_fraction)
    # Validate confidence allocations before mutating any structural state.
    rq = radius(len(x_fresh), delta_q)
    rh = radius(len(x_fresh), delta_h, activation_clip)
    ledger.apply()
    pre = snapshot(model, ledger, "controller_pre_update")
    record = assess(
        model, x_fresh, y_fresh, beta, sigma, noise_seed,
        activation_clip
    )
    lo = max(0.0, record["q_hat"] - rq)
    hi = min(1.0, record["q_hat"] + rq)
    if hi < q_star:
        decision = "certified_crossing"
    elif lo >= q_star:
        decision = "certified_noncrossing"
    else:
        decision = "unresolved"
    severity = min(2.0, max(
        0.5, model.effective_gain()
        + 2.0 * (record["h_hat_clipped"] - 0.1)
    ))
    fraction = min(0.20, max(0.02, base_fraction * severity))
    removed = {n: 0 for n in ledger.masks}
    if decision == "certified_crossing":
        removed = ledger.delete_fraction(
            fraction, permanent=True, minimum_one=minimum_one
        )
    return {
        "pre_update": pre,
        "pre_update_observation": record,
        "q_interval": (lo, hi),
        "activation_radius": rh,
        "decision": decision,
        "unresolved_policy": "no_deletion",
        "severity_rule": "source_form_with_clipped_mean",
        "requested_fraction": fraction,
        "minimum_one_rule": bool(minimum_one),
        "removed_by_layer": removed,
        "post_update": snapshot(model, ledger, "controller_post_update"),
    }

def paired_assessment(model0, model1, x_fresh, y_fresh,
                      beta, sigma, noise_seed, delta,
                      activation_clip=10.0):
    r = radius(len(x_fresh), delta, interval_length=2.0)
    a = assess(model0, x_fresh, y_fresh, beta, sigma,
               noise_seed, activation_clip)
    b = assess(model1, x_fresh, y_fresh, beta, sigma,
               noise_seed, activation_clip)
    difference = a["correctness"] - b["correctness"]
    estimate = float(difference.mean())
    return {
        "contrast": estimate, "radius": r,
        "lower": max(-1.0, estimate - r),
        "upper": min(1.0, estimate + r),
        "challenge": (beta, sigma),
    }