Listing: Listing

import math
from fractions import Fraction
import torch

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

def _ceil_fraction(q):
    return (q.numerator + q.denominator - 1) // q.denominator

@torch.no_grad()
def _install_exact_counts(manager, counts):
    # Validate every layer before changing any mask.
    for name, k in counts.items():
        p = manager.model.get_parameter(name)
        old = manager.masks[name].reshape(-1).bool()
        n_active = int(old.sum().item())
        if not 0 <= int(k) <= n_active:
            raise ValueError(
                f"{name}: target count {k} exceeds active "
                f"count {n_active}; mask regrowth is forbidden"
            )
        if not bool(torch.isfinite(p).all()):
            raise ValueError(f"{name}: non-finite weights")

    stats = {}
    for name, k in counts.items():
        p = manager.model.get_parameter(name)
        old = manager.masks[name].reshape(-1).bool()

        # nonzero returns positions in ascending flattened order.
        eligible = torch.nonzero(old, as_tuple=False).flatten()
        magnitudes = p.detach().reshape(-1)[eligible].abs()
        order = torch.argsort(
            magnitudes, descending=True, stable=True
        )

        new = torch.zeros_like(old)
        if int(k) > 0:
            new[eligible[order[:int(k)]]] = True

        n_before = int(old.sum().item())
        n_after = int(new.sum().item())
        assert n_after == int(k)
        assert not bool((new & ~old).any())

        manager.masks[name] = new.reshape_as(p).to(p.dtype)
        p.mul_(manager.masks[name])

        stats[name] = {
            "kept": n_after,
            "total": p.numel(),
            "actual_sparsity": 1.0 - n_after / p.numel(),
            "removed": n_before - n_after,
            "remaining": n_after,
            "fraction": (
                (n_before - n_after) / n_before
                if n_before else 0.0
            ),
        }
    return stats

def prune_to_density_exact(self, target_density, per_layer=True):
    if not per_layer:
        raise NotImplementedError(
            "This patch implements layer-wise allocation only."
        )
    d = _fraction01(target_density, "target_density")
    counts = {
        name: _ceil_fraction(d * mask.numel())
        for name, mask in self.masks.items()
    }
    stats = _install_exact_counts(self, counts)
    self.history.append(("exact_density", float(d), stats))
    return stats

def prune_by_magnitude_exact(self, sparsity, per_layer=True):
    s = _fraction01(sparsity, "sparsity")
    return prune_to_density_exact(
        self, float(1 - s), per_layer=per_layer
    )

def prune_fraction_of_remaining_exact(
    self, fraction_to_remove, per_layer=True
):
    if not per_layer:
        raise NotImplementedError(
            "This patch implements layer-wise allocation only."
        )
    f = _fraction01(fraction_to_remove, "fraction_to_remove")
    counts = {}
    for name, mask in self.masks.items():
        n = int(mask.bool().sum().item())
        q = f * n
        removed = q.numerator // q.denominator
        counts[name] = n - removed

    stats = _install_exact_counts(self, counts)
    self.history.append(("exact_fraction", float(f), stats))
    return stats

# Install after defining the original manager class.
TaskGatedPruningManager.prune_to_density = (
    prune_to_density_exact
)
TaskGatedPruningManager.prune_by_magnitude = (
    prune_by_magnitude_exact
)
TaskGatedPruningManager.prune_fraction_of_remaining = (
    prune_fraction_of_remaining_exact
)

def structural_density(manager):
    kept = sum(
        int(mask.bool().sum().item())
        for mask in manager.masks.values()
    )
    total = sum(mask.numel() for mask in manager.masks.values())
    return kept / total if total else 0.0

def test_exact_pruning_patch():
    # The original constructor accepts this simple nn.Module
    # when prune_head=True: all matrix weights are eligible.
    model = torch.nn.Linear(4, 1, bias=False)
    with torch.no_grad():
        model.weight.copy_(
            torch.tensor([[3.0, 3.0, 1.0, 0.0]])
        )
    mgr = TaskGatedPruningManager(model)

    mgr.prune_to_density(0.5)
    assert int(mgr.masks["weight"].sum().item()) == 2

    # Stable tie break retains flattened index zero.
    mgr.prune_fraction_of_remaining(0.5)
    expected = torch.tensor([[1.0, 0.0, 0.0, 0.0]])
    assert torch.equal(mgr.masks["weight"], expected)

    try:
        mgr.prune_to_density(0.75)
    except ValueError:
        pass
    else:
        raise AssertionError("Mask regrowth was not rejected")

    # Endpoint behavior: remove every active position.
    mgr.prune_fraction_of_remaining(1.0)
    assert structural_density(mgr) == 0.0
    assert int(torch.count_nonzero(model.weight).item()) == 0