inference_code/inference/pruning/temporal.py

"""Timestep-bin TENP aggregation. Default apply path still sums all bins (plain TENP)."""

from __future__ import annotations

import torch

from .scoring import fuse_modality_scores

DEFAULT_NUM_BINS = 5
DEFAULT_N_MIN = 32
DEFAULT_TEMPORAL_BETA = 0.3
DEFAULT_TEMPORAL_TOPK = 2
TEMPORAL_MODES = ("none", "mean", "ta", "st")
T_MAX_DEFAULT = 1000.0


def raw_signed_projection_score(table: torch.Tensor) -> torch.Tensor:
    """Sum routed signed projections without z-score changing their sign.

    Accepts either ``[3, d_int]`` or ``[K, 3, d_int]`` and sums only the
    modality axis. This score is deliberately separate from the normalized
    TENP ranking score.
    """
    if table.ndim == 2 and table.shape[0] == 3:
        return table.detach().float().sum(dim=0)
    if table.ndim == 3 and table.shape[1] == 3:
        return table.detach().float().sum(dim=1)
    raise ValueError(f"projection table must be [3,d] or [K,3,d], got {tuple(table.shape)}")


def _positive_percentile_rank(
    importance_bin: torch.Tensor,
    signed_bin: torch.Tensor,
) -> torch.Tensor:
    """Rank neurons within each bin, while raw non-positive neurons score zero."""
    if importance_bin.ndim != 2 or signed_bin.shape != importance_bin.shape:
        raise ValueError(
            f"importance/signed bins must share [K,d], got "
            f"{tuple(importance_bin.shape)} and {tuple(signed_bin.shape)}"
        )
    ranks = torch.zeros_like(importance_bin, dtype=torch.float32)
    for bin_id in range(int(importance_bin.shape[0])):
        positive = torch.nonzero(signed_bin[bin_id] > 0, as_tuple=False).reshape(-1)
        n_pos = int(positive.numel())
        if n_pos == 0:
            continue
        order = torch.argsort(importance_bin[bin_id, positive], stable=True)
        percentile = torch.empty(n_pos, device=importance_bin.device, dtype=torch.float32)
        percentile[order] = (
            torch.arange(1, n_pos + 1, device=importance_bin.device, dtype=torch.float32)
            / float(n_pos)
        )
        ranks[bin_id, positive] = percentile
    return ranks


def temporal_criticality(
    importance_bin: torch.Tensor,
    signed_bin: torch.Tensor,
    topq: int = DEFAULT_TEMPORAL_TOPK,
    adjacent: bool = False,
) -> torch.Tensor:
    """Top-Q mean of per-bin positive percentile ranks for each neuron."""
    ranks = _positive_percentile_rank(importance_bin, signed_bin)
    k = int(ranks.shape[0])
    q = min(max(int(topq), 1), k)
    if not adjacent or q == 1:
        return torch.topk(ranks, q, dim=0).values.mean(dim=0)
    windows = torch.stack(
        [ranks[start : start + q].mean(dim=0) for start in range(k - q + 1)],
        dim=0,
    )
    return windows.max(dim=0).values


def temporal_coverage(
    importance_bin: torch.Tensor,
    keep_idx: torch.Tensor,
    eps: float = 1e-8,
) -> torch.Tensor:
    """Fraction of positive per-bin importance retained by ``keep_idx``."""
    if importance_bin.ndim != 2:
        raise ValueError(f"importance_bin must be [K,d], got {tuple(importance_bin.shape)}")
    keep = keep_idx.detach().long().reshape(-1)
    positive = importance_bin.detach().float().clamp_min(0)
    retained = positive[:, keep].sum(dim=1)
    return retained / positive.sum(dim=1).clamp_min(eps)


def timestep_to_bin(
    timestep: float | None,
    num_bins: int,
    step_index: int | None = None,
    num_steps: int | None = None,
    t_max: float = T_MAX_DEFAULT,
) -> int:
    """Map a diffusion step onto `num_bins` equal bins.

    Prefer inference `step_index` so each bin gets a similar number of solver
    steps. Fall back to the scheduler value in `[0, t_max)`. Missing timestep
    lands in bin 0, which keeps legacy dumps identical to plain TENP.
    """
    k = max(int(num_bins), 1)
    if k == 1:
        return 0
    if step_index is not None and num_steps:
        n = max(int(num_steps), 1)
        return min(k - 1, max(0, int(step_index) * k // n))
    if timestep is None:
        return 0
    t = float(timestep)
    if t_max <= 0:
        return 0
    return min(k - 1, max(0, int(t / float(t_max) * k)))


def merge_sparse_bins(
    n_tokens_bin: torch.Tensor,
    n_min: int = DEFAULT_N_MIN,
) -> torch.Tensor:
    """Keep-mask over bins. Sparse bins are folded into the nearest kept neighbor."""
    counts = n_tokens_bin.detach().long().reshape(-1)
    k = int(counts.numel())
    keep = counts >= int(n_min)
    if bool(keep.any()) is False:
        keep = counts > 0
    if bool(keep.any()) is False:
        keep = torch.ones(k, dtype=torch.bool, device=counts.device)
        keep[1:] = False
    return keep


def fold_sparse_bin_mass(
    n_tokens_bin: torch.Tensor,
    values: torch.Tensor | None = None,
    n_min: int = DEFAULT_N_MIN,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Move sparse-bin token mass (and optional values) onto the nearest kept bin.

    Returns (folded_counts, folded_values). `values` is typically D_{e,k};
    when omitted, only counts are folded.
    """
    n = n_tokens_bin.detach().float().reshape(-1).clone()
    keep = merge_sparse_bins(n, n_min=n_min)
    if values is None:
        v = torch.zeros_like(n)
        has_v = False
    else:
        v = values.detach().float().reshape(-1).clone()
        if v.numel() != n.numel():
            raise ValueError(f"values {tuple(v.shape)} != n_tokens_bin {tuple(n.shape)}")
        has_v = True
    kept_idx = torch.nonzero(keep, as_tuple=False).reshape(-1)
    for i in range(int(n.numel())):
        if bool(keep[i]):
            continue
        if float(n[i]) <= 0:
            continue
        nearest = int(kept_idx[torch.argmin((kept_idx - i).abs())].item())
        if has_v:
            mass = float(n[nearest]) + float(n[i])
            if mass > 0:
                v[nearest] = (v[nearest] * n[nearest] + v[i] * n[i]) / mass
            v[i] = 0
        n[nearest] = n[nearest] + n[i]
        n[i] = 0
    return n, v


def routing_sensitivity_alpha(
    n_tokens_bin: torch.Tensor,
    sensitivity: torch.Tensor,
    n_min: int = DEFAULT_N_MIN,
    eps: float = 1e-8,
    layer_sensitivity: torch.Tensor | None = None,
) -> torch.Tensor:
    """α_{e,k} = π_{e,k} D_{e,k} / sum_j(π D). Sparse bins fold into the nearest kept bin."""
    n = n_tokens_bin.detach().float().reshape(-1)
    d = sensitivity.detach().float().reshape(-1)
    if d.numel() != n.numel():
        raise ValueError(f"sensitivity {tuple(d.shape)} != n_tokens_bin {tuple(n.shape)}")
    if layer_sensitivity is not None:
        d_layer = layer_sensitivity.detach().float().reshape(-1)
        if d_layer.numel() == d.numel():
            keep = merge_sparse_bins(n, n_min=n_min)
            d = torch.where(~keep, d_layer, d)
    n, d = fold_sparse_bin_mass(n, d, n_min=n_min)
    pi = n / n.sum().clamp_min(1.0)
    raw = pi * d.clamp_min(0)
    if float(raw.sum()) <= 0:
        raw = pi
    return raw / raw.sum().clamp_min(eps)


def mean_alpha(n_tokens_bin: torch.Tensor, n_min: int = DEFAULT_N_MIN) -> torch.Tensor:
    keep = merge_sparse_bins(n_tokens_bin, n_min=n_min)
    n_fold, _ = fold_sparse_bin_mass(n_tokens_bin, n_min=n_min)
    w = torch.where(n_fold > 0, torch.ones_like(n_fold), torch.zeros_like(n_fold))
    if float(w.sum()) <= 0:
        w = keep.float()
    return w / w.sum().clamp_min(1.0)


def fuse_bins_to_importance(
    table_bin: torch.Tensor,
    n_tokens_bin: torch.Tensor,
    lambdas: tuple[float, float, float],
    n_min: int = DEFAULT_N_MIN,
    zscore: bool = True,
) -> torch.Tensor:
    """Per-bin modality fusion. `table_bin`: [K, 3, d_int] -> [K, d_int]."""
    if table_bin.ndim != 3 or table_bin.shape[1] != 3:
        raise ValueError(f"table_bin must be [K, 3, d_int], got {tuple(table_bin.shape)}")
    k, _, d_int = table_bin.shape
    folded = table_bin.detach().float().clone()
    keep = merge_sparse_bins(n_tokens_bin, n_min=n_min)
    kept_idx = torch.nonzero(keep, as_tuple=False).reshape(-1)
    for bin_id in range(k):
        if bool(keep[bin_id]):
            continue
        if float(folded[bin_id].abs().sum()) == 0:
            continue
        nearest = int(kept_idx[torch.argmin((kept_idx - bin_id).abs())].item())
        folded[nearest] = folded[nearest] + folded[bin_id]
        folded[bin_id] = 0
    out = torch.zeros(k, d_int, device=table_bin.device, dtype=torch.float32)
    for bin_id in range(k):
        if not bool(keep[bin_id]):
            continue
        row = folded[bin_id]
        if float(row.abs().sum()) == 0:
            continue
        out[bin_id] = fuse_modality_scores(row, lambdas=lambdas, zscore=zscore)
    return out


def aggregate_temporal(
    importance_bin: torch.Tensor,
    alpha: torch.Tensor | None = None,
    mode: str = "st",
    beta: float = DEFAULT_TEMPORAL_BETA,
    temporal_topk: int = DEFAULT_TEMPORAL_TOPK,
) -> torch.Tensor:
    """Reduce [K, d_int] bin scores to one neuron vector."""
    if mode not in TEMPORAL_MODES:
        raise ValueError(f"unknown temporal mode {mode!r}; expected one of {TEMPORAL_MODES}")
    if importance_bin.ndim != 2:
        raise ValueError(f"importance_bin must be [K, d_int], got {tuple(importance_bin.shape)}")
    k = int(importance_bin.shape[0])
    if mode == "none":
        raise ValueError("mode 'none' is handled by the summed projection table, not here")
    live = importance_bin.abs().sum(dim=1) > 0
    if not bool(live.any()):
        return importance_bin.sum(dim=0)
    stacked = importance_bin[live]
    if mode == "mean":
        if alpha is None:
            return stacked.mean(dim=0)
        w = alpha.detach().float().reshape(-1)[live]
        w = w / w.sum().clamp_min(1e-8)
        return (w.unsqueeze(-1) * stacked).sum(dim=0)
    if mode == "ta":
        mean = stacked.mean(dim=0)
        topk = min(max(int(temporal_topk), 1), int(stacked.shape[0]))
        peak = stacked.topk(topk, dim=0).values.mean(dim=0)
        return mean + float(beta) * peak
    if alpha is None:
        raise ValueError("ST-TENP needs alpha")
    w = alpha.detach().float().reshape(-1)
    if w.numel() != k:
        raise ValueError(f"alpha {tuple(w.shape)} != K={k}")
    w = torch.where(live, w, torch.zeros_like(w))
    w = w / w.sum().clamp_min(1e-8)
    return (w.unsqueeze(-1) * importance_bin).sum(dim=0)


def bin_coefficient_of_variation(importance_bin: torch.Tensor, eps: float = 1e-8) -> torch.Tensor:
    """Per-neuron CV across timestep bins. Empty bins are ignored."""
    live = importance_bin.abs().sum(dim=1) > 0
    stacked = importance_bin[live] if bool(live.any()) else importance_bin
    mu = stacked.mean(dim=0)
    sigma = stacked.std(dim=0, unbiased=False)
    return sigma / (mu.abs() + eps)


def topk_jaccard(importance_bin: torch.Tensor, frac: float = 0.20) -> torch.Tensor:
    """Pairwise Jaccard of per-bin top-(frac) neurons. Diagonal is 1."""
    k, d_int = importance_bin.shape
    n_keep = max(1, int(round(float(frac) * d_int)))
    n_keep = min(n_keep, d_int)
    sets: list[set[int]] = []
    live = importance_bin.abs().sum(dim=1) > 0
    for bin_id in range(k):
        if not bool(live[bin_id]):
            sets.append(set())
            continue
        _, idx = torch.topk(importance_bin[bin_id], n_keep)
        sets.append(set(int(x) for x in idx.tolist()))
    jac = torch.zeros(k, k, dtype=torch.float32)
    for i in range(k):
        for j in range(k):
            if not sets[i] or not sets[j]:
                jac[i, j] = 0.0 if i != j else 1.0
                continue
            inter = len(sets[i] & sets[j])
            union = len(sets[i] | sets[j])
            jac[i, j] = float(inter) / float(union) if union else 0.0
    return jac