inference_code/inference/pruning/scoring.py

"""Fuse per-modality TENP / energy tables. Default λ = (3, 1, 1)."""

from __future__ import annotations

import torch

from .surgery import round_to_multiple

VIDEO, AUDIO, TEXT = 0, 1, 2
MODALITY_FUSE_MODES = (
    "fixed133",
    "fixed211",
    "fixed311",
    "fixed511",
    "fixed3011",
    "equal",
    "pooled",
    "invfreq",
    "invsqrt",
    "invlog",
    # rank<λVλAλT>: per-modality percentile-rank normalisation instead of z-score.
    # z-score on a heavy-tailed projection table is dominated by a few extreme neurons
    # (kurtosis 1e2..1e3), so λ only reorders near-zero neurons; ranks spread the bulk.
    "rank111",
    "rank133",
    "rank211",
    "rank311",
    "rank511",
    "rank3011",
)
DEFAULT_MODALITY_FUSE = "fixed311"
DEFAULT_MODALITY_LAMBDAS = (3.0, 1.0, 1.0)
_FIXED_LAMBDAS = {
    "111": (1.0, 1.0, 1.0),
    "133": (1.0, 3.0, 3.0),
    "211": (2.0, 1.0, 1.0),
    "311": (3.0, 1.0, 1.0),
    "511": (5.0, 1.0, 1.0),
    "3011": (30.0, 1.0, 1.0),
}
# No per-modality z-score: λ is applied to the raw SUM table C_m (paper Eq. modality-contribution).
# pooled = λ=(1,1,1) on sums, so video token count dominates.
# invfreq / invsqrt / invlog set λ_m from routed N_m so the fusion ratio is not 1:1:1.
NO_ZSCORE_FUSE_MODES = frozenset({"pooled", "sum", "nozscore", "invfreq", "invsqrt", "invlog"})
POOLED_FUSE_MODES = NO_ZSCORE_FUSE_MODES
FUSE_NORMS = ("z", "sum", "rank")


def fuse_norm(mode: str) -> str:
    """Per-modality normalisation for a fuse mode: 'z' (default), 'sum' (none) or 'rank'."""
    key = str(mode).strip().lower()
    if key in NO_ZSCORE_FUSE_MODES:
        return "sum"
    if key.startswith("rank"):
        return "rank"
    return "z"


def fuse_zscore(mode: str) -> bool:
    """False for count-weighted SUM fusion (no per-modality z-score)."""
    return fuse_norm(mode) == "z"


def rank_normalize(row: torch.Tensor) -> torch.Tensor:
    """Centered percentile rank in [-0.5, 0.5]; ties broken by index, scale-free."""
    d = int(row.numel())
    if d < 2:
        return torch.zeros_like(row, dtype=torch.float32)
    order = torch.argsort(row.float(), stable=True)
    ranks = torch.empty(d, dtype=torch.float32, device=row.device)
    ranks[order] = torch.arange(d, dtype=torch.float32, device=row.device)
    return ranks / float(d - 1) - 0.5


def _token_count_vector(
    n_tokens_mod: torch.Tensor | list[int] | tuple[int, ...] | None,
    mode: str,
) -> torch.Tensor:
    if n_tokens_mod is None:
        raise ValueError(f"{mode} fuse requires n_tokens_mod")
    n = torch.as_tensor(n_tokens_mod, dtype=torch.float64).reshape(-1)
    if int(n.numel()) < 3:
        raise ValueError(f"{mode} fuse expected 3 modality counts, got {tuple(n.shape)}")
    return n[:3].clamp_min(1.0)


def modality_lambdas(
    n_tokens_mod: torch.Tensor | list[int] | tuple[int, ...] | None,
    mode: str = DEFAULT_MODALITY_FUSE,
) -> tuple[float, float, float]:
    key = str(mode).strip().lower()
    if key in ("pooled", "sum", "nozscore"):
        return (1.0, 1.0, 1.0)
    if key in ("invfreq", "invsqrt", "invlog"):
        n = _token_count_vector(n_tokens_mod, key)
        if key == "invfreq":
            w = 1.0 / n
        elif key == "invsqrt":
            w = 1.0 / n.sqrt()
        else:
            w = 1.0 / (1.0 + n).log()
        return tuple(float(x) for x in w.tolist())
    if key not in MODALITY_FUSE_MODES:
        raise ValueError(f"unknown modality fuse {mode!r}; expected one of {MODALITY_FUSE_MODES}")
    if key == "equal":
        return (1.0, 1.0, 1.0)
    for prefix in ("fixed", "rank"):
        if key.startswith(prefix) and key[len(prefix):] in _FIXED_LAMBDAS:
            return _FIXED_LAMBDAS[key[len(prefix):]]
    return DEFAULT_MODALITY_LAMBDAS


def fuse_modality_scores(
    table: torch.Tensor,
    lambdas: tuple[float, float, float] = DEFAULT_MODALITY_LAMBDAS,
    eps: float = 1e-6,
    zscore: bool = True,
    norm: str | None = None,
) -> torch.Tensor:
    """table: [3, d_int] for V/A/T. Empty modalities are skipped, not imputed.

    `norm` overrides `zscore`: 'z' (z-score), 'sum' (raw), 'rank' (centered percentile rank).
    """
    if table.ndim != 2 or table.shape[0] != 3:
        raise ValueError(f"table must be [3, d_int], got {tuple(table.shape)}")
    if norm is None:
        norm = "z" if zscore else "sum"
    if norm not in FUSE_NORMS:
        raise ValueError(f"unknown norm {norm!r}; expected one of {FUSE_NORMS}")
    parts = []
    for modality, weight in enumerate(lambdas):
        row = table[modality].float()
        if float(row.abs().sum()) == 0:
            continue
        if norm == "z":
            mu = row.mean()
            sigma = row.std(unbiased=False).clamp_min(eps)
            parts.append(weight * (row - mu) / sigma)
        elif norm == "rank":
            parts.append(weight * rank_normalize(row))
        else:
            parts.append(weight * row)
    if not parts:
        raise ValueError("no routed tokens in any modality")
    return sum(parts)


def keep_indices(score: torch.Tensor, p_mid: float, multiple: int = 4) -> torch.Tensor:
    d_int = score.numel()
    n_keep = max(multiple, round_to_multiple(int((1.0 - p_mid) * d_int), multiple))
    n_keep = min(n_keep, d_int)
    _, idx = torch.topk(score, n_keep)
    idx, _ = torch.sort(idx)
    return idx


def random_keep_indices(
    d_int: int,
    p_mid: float,
    generator: torch.Generator | None = None,
    multiple: int = 4,
) -> torch.Tensor:
    n_keep = max(multiple, round_to_multiple(int((1.0 - p_mid) * d_int), multiple))
    n_keep = min(n_keep, d_int)
    idx = torch.randperm(d_int, generator=generator)[:n_keep]
    idx, _ = torch.sort(idx)
    return idx


def fuse_global_modality_scores(
    table: torch.Tensor,
    lambdas: tuple[float, float, float] = DEFAULT_MODALITY_LAMBDAS,
    dead_mask: torch.Tensor | None = None,
    magnitude: torch.Tensor | None = None,
    eps: float = 1e-6,
    zscore: bool = True,
) -> torch.Tensor:
    """table: [n_mlp, 3, d_int]. Z-score each modality over live MLPs, then fuse."""
    if table.ndim != 3 or table.shape[1] != 3:
        raise ValueError(f"table must be [n_mlp, 3, d_int], got {tuple(table.shape)}")
    n_mlp, _, d_int = table.shape
    live = (
        torch.ones(n_mlp, dtype=torch.bool, device=table.device)
        if dead_mask is None
        else ~dead_mask.bool()
    )
    score = torch.zeros(n_mlp, d_int, device=table.device, dtype=torch.float32)
    parts = 0
    for modality, weight in enumerate(lambdas):
        row = table[:, modality].float()
        vals = row[live]
        if vals.numel() == 0 or float(vals.abs().sum()) == 0:
            continue
        if zscore:
            mu = vals.mean()
            sigma = vals.std(unbiased=False).clamp_min(eps)
            z = (row - mu) / sigma
        else:
            z = row
        z[~live] = 0
        score = score + weight * z
        parts += 1
    if parts == 0 and (dead_mask is None or not bool(dead_mask.all())):
        raise ValueError("no routed tokens in any modality")
    if dead_mask is not None and bool(dead_mask.any()):
        dead = dead_mask.bool()
        tie = torch.zeros(n_mlp, d_int, device=table.device, dtype=torch.float32)
        if magnitude is not None:
            mag = magnitude.float()
            live_mag = mag[dead]
            sigma = live_mag.std(unbiased=False).clamp_min(eps)
            tie[dead] = (mag[dead] - live_mag.mean()) / sigma
        score[dead] = -1e6 + 1e-3 * tie[dead]
    return score


def allocate_global_keeps(
    score: torch.Tensor,
    p_mid: float,
    multiple: int = 4,
    min_keep: int = 4,
) -> list[torch.Tensor]:
    """Global top-(1-p_mid) over all MLPs, then round each MLP's keep to `multiple`."""
    if score.ndim != 2:
        raise ValueError(f"score must be [n_mlp, d_int], got {tuple(score.shape)}")
    n_mlp, d_int = score.shape
    min_keep = min(max(multiple, min_keep), d_int)
    n_keep = max(
        n_mlp * min_keep,
        round_to_multiple(int((1.0 - p_mid) * n_mlp * d_int), multiple),
    )
    n_keep = min(n_keep, n_mlp * d_int)
    selected = torch.zeros(n_mlp, d_int, dtype=torch.bool, device=score.device)
    _, top = torch.topk(score.reshape(-1), n_keep)
    selected.view(-1)[top] = True
    keeps = []
    for i in range(n_mlp):
        idx = selected[i].nonzero(as_tuple=False).reshape(-1)
        n = int(idx.numel())
        n_target = min(d_int, max(min_keep, round_to_multiple(n if n else min_keep, multiple)))
        if n_target == n:
            keep = idx
        elif n_target < n:
            _, order = torch.topk(score[i, idx], n_target)
            keep = idx[order]
        else:
            extra = n_target - n
            masked = score[i].clone()
            if n:
                masked[idx] = -float("inf")
            _, add = torch.topk(masked, extra)
            keep = torch.cat([idx, add]) if n else add
        keep, _ = torch.sort(keep)
        keeps.append(keep)
    return keeps