Listing: Listing

import copy
import hashlib
import math
from fractions import Fraction
from collections import Counter, defaultdict
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 frac01(value):
    value = finite(value, "fraction")
    if not 0 <= value <= 1:
        raise ValueError("fraction outside [0, 1]")
    return Fraction(str(value))

def generator(seed, role):
    # Stable role separation: Python's randomized hash() is not used.
    message = f"{int(seed)}:{role}".encode("utf-8")
    digest = hashlib.sha256(message).digest()
    value = int.from_bytes(digest[:8], "little") % (2**63 - 1)
    return torch.Generator(device="cpu").manual_seed(value)

def noise_streams(seed):
    return {
        name: generator(seed, name)
        for name in ("external_input", "ordinary_input", "biased_hidden")
    }

class CheckedGRUNetwork(nn.Module):
    def __init__(self, hidden_size=384, seed=0, dtype=torch.float64):
        super().__init__()
        if not isinstance(hidden_size, int) or hidden_size <= 0:
            raise ValueError("positive integer hidden size required")
        self.hidden_size = hidden_size
        # Only the CPU default generator is changed inside this context.
        with torch.random.fork_rng(devices=[]):
            torch.default_generator.manual_seed(int(seed))
            self.gru_cell = nn.GRUCell(2, hidden_size, dtype=dtype)
            self.fc = nn.Linear(hidden_size, 4, dtype=dtype)
        for tag, parameter in self.weight_items():
            self.register_buffer(
                "M_" + tag, torch.ones_like(parameter, dtype=torch.bool)
            )
            self.register_buffer(
                "S_" + tag, torch.zeros_like(parameter, dtype=torch.bool)
            )
        self.a, self.g, self.p = 1.0, 1.0, 0.0
        self.activation = "relu"
        self.decay_rate = 0.0

    def weight_items(self):
        return (
            ("ih", self.gru_cell.weight_ih),
            ("hh", self.gru_cell.weight_hh),
            ("out", self.fc.weight),
        )

    def effective_gain(self):
        a = finite(self.a, "a")
        g = finite(self.g, "g")
        p = 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 validate_structure(self):
        for tag, parameter in self.weight_items():
            mask = getattr(self, "M_" + tag)
            scar = getattr(self, "S_" + tag)
            if mask.dtype != torch.bool or scar.dtype != torch.bool:
                raise ValueError("binary boolean structural state required")
            if mask.shape != parameter.shape or scar.shape != parameter.shape:
                raise ValueError("structural shape mismatch")
            if bool((mask & scar).any()):
                raise ValueError("active coordinate intersects scar")
            if not bool(torch.isfinite(parameter).all()):
                raise ValueError("non-finite weight")

    @torch.no_grad()
    def enforce(self):
        self.validate_structure()
        for tag, parameter in self.weight_items():
            parameter.mul_(getattr(self, "M_" + tag).to(parameter.dtype))

    @torch.no_grad()
    def delete_fraction(self, fraction, permanent=False):
        f = frac01(fraction)
        self.validate_structure()
        plans = {}
        for tag, parameter in self.weight_items():
            mask = getattr(self, "M_" + tag)
            eligible = torch.nonzero(mask.reshape(-1)).flatten()
            amount = f * len(eligible)
            k = amount.numerator // amount.denominator
            values = parameter.detach().reshape(-1)[eligible].abs()
            order = torch.argsort(values, stable=True)
            plans[tag] = eligible[order[:k]]
        for tag, selected in plans.items():
            getattr(self, "M_" + tag).view(-1)[selected] = False
            if permanent:
                getattr(self, "S_" + tag).view(-1)[selected] = True
        self.enforce()
        return {tag: len(selected) for tag, selected in plans.items()}

    def forward(self, sequence, mode="ordinary", xi=0.0, tau=0.0,
                sigma=0.0, beta=0.0, streams=None, return_states=False):
        if sequence.device.type != "cpu":
            raise ValueError("this repair core is CPU-specific")
        if sequence.ndim != 3 or sequence.shape[1] < 1 \
                or sequence.shape[2] != 2:
            raise ValueError("expected nonempty [N, T, 2] sequence")
        if sequence.dtype != self.gru_cell.weight_ih.dtype:
            raise ValueError("input and model dtypes must agree")
        if not bool(torch.isfinite(sequence).all()):
            raise ValueError("non-finite sequence")
        xi, tau = finite(xi, "xi"), finite(tau, "tau")
        sigma, beta = finite(sigma, "sigma"), finite(beta, "beta")
        if min(xi, tau, sigma) < 0:
            raise ValueError("negative noise intensity")
        if mode not in ("ordinary", "biased"):
            raise ValueError("explicit mode required")
        if mode == "ordinary" and (sigma != 0 or beta != 0):
            raise ValueError("ordinary mode has no hidden noise or bias")
        if mode == "biased" and tau != 0:
            raise ValueError("biased mode has no ordinary input stress")
        if self.activation not in ("relu", "tanh"):
            raise ValueError("unsupported external activation")
        ge = self.effective_gain()
        phi = torch.relu if self.activation == "relu" else torch.tanh
        h = sequence.new_zeros((len(sequence), self.hidden_size))
        states = []

        def draw(reference, role):
            if streams is None or role not in streams:
                raise ValueError("missing explicit random stream")
            return torch.randn(
                reference.shape, dtype=reference.dtype,
                generator=streams[role]
            )

        for t in range(sequence.shape[1]):
            inp = sequence[:, t, :]
            if xi > 0:
                inp = inp + xi * draw(inp, "external_input")
            inp = ge * inp
            if mode == "ordinary" and tau > 0:
                inp = inp + tau * draw(inp, "ordinary_input")
            h = self.gru_cell(inp, h)
            if mode == "biased":
                if sigma > 0:
                    h = h + sigma * draw(h, "biased_hidden")
                h = h + beta * (1.0 - 0.3 * self.p)
            h = phi(self.a * h - 0.15 * self.p)
            if return_states:
                states.append(h)
        output = self.fc(h)
        return (output, torch.stack(states, dim=1)) \
            if return_states else output

def phase_snapshot(model, phase):
    model.validate_structure()
    return {
        "phase": str(phase),
        "hidden_size": model.hidden_size,
        "dtype": model.gru_cell.weight_ih.dtype,
        "state_dict": {
            key: value.detach().clone()
            for key, value in model.state_dict().items()
        },
        "controls": {
            "a": model.a, "g": model.g, "p": model.p,
            "activation": model.activation,
            "decay_rate": model.decay_rate,
        },
        "training": model.training,
        "optimizer_policy": "fresh_optimizer_at_next_training_call",
    }

def restore_phase(saved):
    model = CheckedGRUNetwork(
        saved["hidden_size"], seed=0, dtype=saved["dtype"]
    )
    model.load_state_dict(saved["state_dict"])
    for key, value in saved["controls"].items():
        setattr(model, key, value)
    model.train(saved["training"])
    model.enforce()
    return model

@torch.no_grad()
def evaluate_explicit(model, sequence, labels, seed, batch_size=1000,
                      mode="ordinary", xi=0.0, tau=0.0,
                      sigma=0.0, beta=0.0):
    if len(sequence) == 0 or len(sequence) != len(labels):
        raise ValueError("nonempty aligned assessment data required")
    if not isinstance(batch_size, int) or batch_size <= 0:
        raise ValueError("positive integer batch size required")
    streams = noise_streams(seed)
    was_training = model.training
    correctness = []
    try:
        model.eval()
        for start in range(0, len(sequence), batch_size):
            xb = sequence[start:start + batch_size]
            yb = labels[start:start + batch_size]
            logits = model(
                xb, mode=mode, xi=xi, tau=tau,
                sigma=sigma, beta=beta, streams=streams
            )
            correctness.append((logits.argmax(dim=1) == yb).double())
    finally:
        model.train(was_training)
    values = torch.cat(correctness)
    return {
        "accuracy": float(values.mean()),
        "correctness": values,
        "n_examples": len(values),
        "sequence_length": sequence.shape[1],
        "mode": mode,
        "noise": (xi, tau, sigma),
        "bias": beta,
        "seed": int(seed),
        "batch_size": batch_size,
    }

def checked_match(rows, target, tolerance, metric="d_w"):
    target, tolerance = finite(target, "target"), \
                        finite(tolerance, "tolerance")
    if target < 0 or tolerance < 0:
        raise ValueError("invalid matching controls")
    eligible = []
    for row in rows:
        value = finite(row["budget"][metric], metric)
        distance = abs(value - target)
        if distance <= tolerance:
            eligible.append((distance, str(row["candidate_id"]), row))
    if not eligible:
        return None
    return min(eligible, key=lambda x: (x[0], x[1]))[2]