Listing: Listing

class ReferenceMaskedLinear(nn.Module):
    def __init__(self, inputs, outputs, bias=False, dtype=torch.float64):
        super().__init__()
        self.weight = nn.Parameter(torch.zeros(outputs, inputs, dtype=dtype))
        self.bias = nn.Parameter(torch.zeros(outputs, dtype=dtype)) \
            if bias else None
        self.register_buffer("mask", torch.ones_like(self.weight, dtype=torch.bool))
        self.register_buffer("reference", torch.zeros_like(self.weight))
        self.register_buffer("reference_valid", torch.tensor(False))
        self.record_effective = False
        self.effective_uses = []

    @torch.no_grad()
    def capture_dense_reference(self):
        if not bool(self.mask.all()):
            raise ValueError("capture reference before pruning")
        if not bool(torch.isfinite(self.weight).all()):
            raise ValueError("non-finite reference")
        self.reference.copy_(self.weight)
        self.reference_valid.fill_(True)

    def begin_scoring(self):
        self.effective_uses.clear()
        self.record_effective = True

    def end_scoring(self):
        self.record_effective = False
        self.effective_uses.clear()

    def forward(self, x):
        effective = self.weight * self.mask.to(self.weight.dtype)
        if self.record_effective and torch.is_grad_enabled():
            effective.retain_grad()
            self.effective_uses.append(effective)
        return F.linear(x, effective, self.bias)

    def effective_gradient(self):
        if not self.effective_uses:
            raise ValueError("no recorded differentiable forward use")
        if any(t.grad is None for t in self.effective_uses):
            raise ValueError("loss must use every recorded forward tensor")
        return sum(
            (t.grad.detach() for t in self.effective_uses),
            torch.zeros_like(self.weight)
        )

    @torch.no_grad()
    def enforce(self):
        self.weight.mul_(self.mask.to(self.weight.dtype))

    @torch.no_grad()
    def delete_exact(self, count):
        active = torch.nonzero(self.mask.reshape(-1)).flatten()
        if not isinstance(count, int) or not 0 <= count <= len(active):
            raise ValueError("invalid exact deletion count")
        values = self.weight.detach().reshape(-1)[active].abs()
        order = torch.argsort(values, stable=True)
        selected = active[order[:count]]
        self.mask.view(-1)[selected] = False
        self.enforce()
        return selected.clone()

    @torch.no_grad()
    def reopen_exact(self, count, score, scale=1.0, positive_only=False):
        if not bool(self.reference_valid):
            raise ValueError("pre-pruning reference is unavailable")
        scale = float(scale)
        if not math.isfinite(scale):
            raise ValueError("finite restoration scale required")
        if score.shape != self.weight.shape \
                or not bool(torch.isfinite(score).all()):
            raise ValueError("finite aligned score required")
        eligible = torch.nonzero((~self.mask).reshape(-1)).flatten()
        values = score.detach().reshape(-1)[eligible].abs()
        if positive_only:
            keep = values > 0
            eligible, values = eligible[keep], values[keep]
        if not isinstance(count, int) or not 0 <= count <= len(eligible):
            raise ValueError("count exceeds declared repair eligibility")
        order = torch.argsort(values, descending=True, stable=True)
        selected = eligible[order[:count]]
        self.mask.view(-1)[selected] = True
        self.weight.view(-1)[selected] = (
            scale * self.reference.view(-1)[selected]
        )
        self.enforce()
        return {
            "selected": selected.clone(),
            "numerically_zero_initializations": int(
                (self.weight.view(-1)[selected] == 0).sum()
            ),
            "score": "effective-weight derivative supplied by caller",
            "initialization": "scaled immutable pre-pruning reference",
        }