Blame for 14B/code/losses/video_nystrom_mmd.py

"""Single-encoder distribution matching: exact repulsion and reference attraction.

Reading map: offline helpers build the persistent Nyström reference; training
starts at ``VideoNystromMMD.forward``. Paper Sec. 3.2 calls its combination
with the Monte Carlo reference minibatch the hybrid reference kernel mean.
The released multi-encoder path uses full-population partial gradients with
FIFO and attraction EMA disabled; those optional modes are retained below.

The released estimator includes diagonal pairs (B * B denominator) and the
reference-reference constant. This differs from the paper's self-excluding
repulsion and constant-free optimization expression; preserve both here,
including their effect on the release's subsequent detached loss scaling.

Every input sample is a complete ``[N, D]`` dense-video feature matrix.  For
generated samples ``G``, frozen landmarks ``Z`` and a frozen reference pack,
the legacy objective is

    mean(K_gg) - 2 * mean(K_gz @ alpha) + k_rr.

An opt-in matched-teacher path additionally accepts a fresh teacher batch
``T``. It computes the complete biased Gaussian-kernel terms ``mean(K_gt)``
and ``mean(K_tt)`` without sampling or approximation. The fixed-reference and
fresh-teacher MMDs are combined with normalized non-negative group weights, so
enabling a second teacher target does not duplicate the student repulsion or
silently double the total attraction coefficient.

The explicit partial-gradient interface uses gradient students ``G`` and a
detached full student population ``S``. Without explicit student indices it
preserves the rectangular estimator ``mean K(G,S)`` and evaluates attractions
only for ``G``. Supplying the positions of ``G`` in ``S`` opts into the
full-student estimator: selected differentiable values are reinserted into
detached ``S``, then the loss uses ``mean K(S,S)``, fixed attraction over
``S``, and matched-teacher attraction over ``S x T``. Thus only ``G`` gets
gradients, including both sides of the full self-kernel. Both partial modes
reject FIFO and attraction EMA because either stateful transform would change
their stated estimators.

By default, ``K_gg`` is the exact, biased Gaussian-RBF Gram matrix (including
its diagonal).  An optional transactional FIFO changes only this repulsion
term to the current-current off-diagonal sum plus the current-history cross
sum, normalized by the exact number of ordered pairs.  FIFO entries are
detached FP32 feature matrices and are published only after optimizer success.
Only the generated samples receive gradients.  The reference attraction is
represented by a kernel-ridge/Nyström fit and ``k_rr`` is the exact biased
reference-reference constant.

An optional transactional path applies an EMA to the complete scalar
attraction term only.  The exact current-batch repulsion is never averaged.
The current attraction keeps its natural ``1-beta`` derivative after the
first committed update; historical attraction values are detached.
"""

from __future__ import annotations

import hashlib
import hmac
import json
import math
import os
import tempfile
from collections.abc import Mapping, Sequence
from typing import Any, Dict, Optional, Tuple

import torch
from torch import nn


REFERENCE_PACK_VERSION = 1
_PACK_KEYS = {
    "format_version",
    "landmarks",
    "alpha",
    "k_rr",
    "sigma",
    "metadata",
    "content_hash",
}


def canonicalize_metadata(metadata: Mapping[str, Any]) -> Dict[str, Any]:
    """Return a deterministic, JSON-safe metadata mapping."""
    if not isinstance(metadata, Mapping):
        raise TypeError("metadata must be a mapping")
    try:
        encoded = json.dumps(
            dict(metadata),
            sort_keys=True,
            separators=(",", ":"),
            ensure_ascii=True,
            allow_nan=False,
        )
    except (TypeError, ValueError) as exc:
        raise ValueError(
            "metadata must contain only finite JSON-serializable values"
        ) from exc
    decoded = json.loads(encoded)
    if not isinstance(decoded, dict):
        raise ValueError("metadata must canonicalize to a dictionary")
    return decoded


def _metadata_contains(actual: Any, expected: Any, path: str = "metadata") -> None:
    """Require ``expected`` to be a recursively exact subset of ``actual``."""
    if isinstance(expected, Mapping):
        if not isinstance(actual, Mapping):
            raise ValueError(f"{path} is not a mapping")
        for key, expected_value in expected.items():
            if key not in actual:
                raise ValueError(f"{path} is missing required key {key!r}")
            _metadata_contains(
                actual[key], expected_value, f"{path}.{key}")
        return
    if actual != expected:
        raise ValueError(
            f"{path} mismatch: saved={actual!r}, expected={expected!r}")


_FINITE_VALIDATION_MAX_ELEMENTS = 64 * 1024 * 1024


def _validate_video_tensor(tensor: torch.Tensor, name: str) -> Tuple[int, int, int]:
    if not isinstance(tensor, torch.Tensor):
        raise TypeError(f"{name} must be a torch.Tensor")
    if tensor.ndim != 3 or min(tensor.shape) <= 0:
        raise ValueError(f"{name} must have non-empty shape [B, N, D]")
    if not tensor.is_floating_point():
        raise TypeError(f"{name} must be floating point")

    # A full isfinite mask is one byte per element.  The production reference
    # has 12,212 * 1,024 * 1,024 elements, so materializing that mask needs
    # 11.93 GiB and can OOM an otherwise valid H200 consolidation.  Validate
    # complete rows in bounded slices instead; every value is still checked.
    row_elements = int(tensor.shape[1]) * int(tensor.shape[2])
    rows_per_chunk = max(
        1, int(_FINITE_VALIDATION_MAX_ELEMENTS) // row_elements)
    detached = tensor.detach()
    for start in range(0, int(tensor.shape[0]), rows_per_chunk):
        block = detached[start:start + rows_per_chunk]
        if not bool(torch.isfinite(block).all().item()):
            raise FloatingPointError(f"{name} contains non-finite values")
    return int(tensor.shape[0]), int(tensor.shape[1]), int(tensor.shape[2])


def _flat_row_norm_sq(flattened: torch.Tensor) -> torch.Tensor:
    """FP32 squared row norms without a full elementwise-square temporary."""
    if flattened.ndim != 2 or not flattened.is_floating_point():
        raise ValueError("flattened video features must be a floating [B, F] tensor")
    return torch.linalg.vector_norm(flattened, ord=2, dim=1).square()


def _pairwise_flat_sq(
    left_flat: torch.Tensor,
    right_flat: torch.Tensor,
    left_sq: torch.Tensor,
    right_sq: torch.Tensor,
) -> torch.Tensor:
    return (
        left_sq[:, None]
        + right_sq[None, :]
        - 2.0 * (left_flat @ right_flat.t())
    ).clamp_min(0)


def _inverse_two_sigma_sq(
    sigma: torch.Tensor | float,
    *,
    device: torch.device,
) -> torch.Tensor:
    sigma_tensor = torch.as_tensor(sigma, device=device, dtype=torch.float32)
    if sigma_tensor.numel() != 1:
        raise ValueError("sigma must be scalar")
    if not bool(torch.isfinite(sigma_tensor).item()) or float(sigma_tensor) <= 0:
        raise ValueError("sigma must be finite and positive")
    return 0.5 / sigma_tensor.square()


def _rbf_flat_precomputed(
    left_flat: torch.Tensor,
    right_flat: torch.Tensor,
    left_sq: torch.Tensor,
    right_sq: torch.Tensor,
    inverse_two_sigma_sq: torch.Tensor,
) -> torch.Tensor:
    distances_sq = _pairwise_flat_sq(
        left_flat, right_flat, left_sq, right_sq)
    return torch.exp(-distances_sq * inverse_two_sigma_sq)


def pairwise_frobenius_sq(
    left: torch.Tensor,
    right: torch.Tensor,
) -> torch.Tensor:
    """Squared Frobenius distances between ``[B,N,D]`` video matrices."""
    _, left_n, left_d = _validate_video_tensor(left, "left")
    _, right_n, right_d = _validate_video_tensor(right, "right")
    if (left_n, left_d) != (right_n, right_d):
        raise ValueError(
            "left/right video matrix shapes differ: "
            f"{(left_n, left_d)} != {(right_n, right_d)}")
    left_flat = left.float().flatten(1)
    right_flat = right.float().flatten(1)
    return _pairwise_flat_sq(
        left_flat,
        right_flat,
        _flat_row_norm_sq(left_flat),
        _flat_row_norm_sq(right_flat),
    )


def rbf_frobenius_kernel(
    left: torch.Tensor,
    right: torch.Tensor,
    sigma: torch.Tensor | float,
) -> torch.Tensor:
    """Gaussian RBF kernel using complete video-matrix samples."""
    _, left_n, left_d = _validate_video_tensor(left, "left")
    _, right_n, right_d = _validate_video_tensor(right, "right")
    if (left_n, left_d) != (right_n, right_d):
        raise ValueError(
            "left/right video matrix shapes differ: "
            f"{(left_n, left_d)} != {(right_n, right_d)}")
    left_flat = left.float().flatten(1)
    right_flat = right.float().flatten(1)
    return _rbf_flat_precomputed(
        left_flat,
        right_flat,
        _flat_row_norm_sq(left_flat),
        _flat_row_norm_sq(right_flat),
        _inverse_two_sigma_sq(sigma, device=left.device),
    )


@torch.no_grad()
def median_frobenius_sigma(
    reference: torch.Tensor,
    *,
    max_samples: Optional[int] = 2048,
    chunk_size: int = 256,
    min_sigma: float = 1.0e-6,
) -> torch.Tensor:
    """Deterministic median-distance bandwidth over video samples.

    A deterministic linspace subset bounds the number of pair distances.
    Pair blocks are constructed in chunks; only the final one-dimensional
    vector of upper-triangular distances is retained for the exact median of
    that subset.
    """
    count, _, _ = _validate_video_tensor(reference, "reference")
    if chunk_size <= 0:
        raise ValueError("chunk_size must be positive")
    if not math.isfinite(float(min_sigma)) or float(min_sigma) <= 0:
        raise ValueError("min_sigma must be finite and positive")
    samples = reference.detach().float()
    if max_samples is not None:
        if int(max_samples) < 2:
            raise ValueError("max_samples must be at least 2 or None")
        if count > int(max_samples):
            indices = torch.linspace(
                0,
                count - 1,
                steps=int(max_samples),
                device=samples.device,
                dtype=torch.float64,
            ).round().long()
            samples = samples.index_select(0, indices)
    count = int(samples.shape[0])
    if count < 2:
        return samples.new_tensor(1.0)

    flattened = samples.flatten(1)
    norms_sq = _flat_row_norm_sq(flattened)
    pieces = []
    for row_start in range(0, count, chunk_size):
        row_end = min(row_start + chunk_size, count)
        for col_start in range(row_start, count, chunk_size):
            col_end = min(col_start + chunk_size, count)
            block = _pairwise_flat_sq(
                flattened[row_start:row_end],
                flattened[col_start:col_end],
                norms_sq[row_start:row_end],
                norms_sq[col_start:col_end],
            )
            if col_start == row_start:
                upper = torch.triu_indices(
                    row_end - row_start,
                    col_end - col_start,
                    offset=1,
                    device=block.device,
                )
                values = block[upper[0], upper[1]]
            else:
                values = block.reshape(-1)
            positive = values[values > 0]
            if positive.numel():
                pieces.append(positive)
    if not pieces:
        return samples.new_tensor(1.0)
    median_sq = torch.cat(pieces).median().clamp_min(float(min_sigma) ** 2)
    return median_sq.sqrt()


def _nearest_centroid_labels(
    samples: torch.Tensor,
    centroids: torch.Tensor,
    chunk_size: int,
) -> Tuple[torch.Tensor, torch.Tensor]:
    sample_flat = samples.float().flatten(1)
    centroid_flat = centroids.float().flatten(1)
    sample_sq = _flat_row_norm_sq(sample_flat)
    centroid_sq = _flat_row_norm_sq(centroid_flat)
    labels = []
    minimum_sq = []
    for start in range(0, samples.shape[0], chunk_size):
        distances = _pairwise_flat_sq(
            sample_flat[start:start + chunk_size],
            centroid_flat,
            sample_sq[start:start + chunk_size],
            centroid_sq,
        )
        values, indices = distances.min(dim=1)
        labels.append(indices)
        minimum_sq.append(values)
    return torch.cat(labels), torch.cat(minimum_sq)


def _deterministic_cluster_sums(
    flattened: torch.Tensor,
    labels: torch.Tensor,
    num_clusters: int,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Segment-sum rows without CUDA atomic ``index_add_``.

    Sorting once and differencing a fixed-order prefix sum makes a rebuilt
    reference pack bitwise reproducible on the same hardware/software stack.
    CUDA ``index_add_`` uses repeated-index atomics whose accumulation order
    is intentionally unspecified.
    """
    counts = torch.bincount(labels, minlength=int(num_clusters))
    order = torch.argsort(labels, stable=True)
    sorted_values = flattened.index_select(0, order)
    prefix = sorted_values.cumsum(dim=0)
    sums = flattened.new_zeros(
        (int(num_clusters), flattened.shape[1]))
    nonempty = counts > 0
    ends = counts.cumsum(0) - 1
    selected_ends = ends[nonempty]
    cluster_sums = prefix.index_select(0, selected_ends)
    starts = selected_ends - counts[nonempty] + 1
    has_predecessor = starts > 0
    if bool(has_predecessor.any()):
        cluster_sums[has_predecessor] -= prefix.index_select(
            0, starts[has_predecessor] - 1)
    sums[nonempty] = cluster_sums
    return sums, counts


@torch.no_grad()
def deterministic_video_kmeans(
    reference: torch.Tensor,
    num_landmarks: int,
    *,
    iterations: int = 20,
    chunk_size: int = 256,
) -> torch.Tensor:
    """Deterministic farthest-first initialized Lloyd k-means."""
    count, num_tokens, feature_dim = _validate_video_tensor(
        reference, "reference")
    num_landmarks = int(num_landmarks)
    if num_landmarks <= 0 or num_landmarks > count:
        raise ValueError(
            f"num_landmarks must be in [1, {count}], got {num_landmarks}")
    if int(iterations) < 0:
        raise ValueError("iterations must be non-negative")
    if int(chunk_size) <= 0:
        raise ValueError("chunk_size must be positive")

    samples = reference.detach().float()
    flattened = samples.flatten(1)
    norms_sq = _flat_row_norm_sq(flattened)
    global_mean = samples.mean(dim=0, keepdim=True).flatten(1)
    mean_sq = _flat_row_norm_sq(global_mean)
    first = int(
        _pairwise_flat_sq(flattened, global_mean, norms_sq, mean_sq)
        .squeeze(1)
        .argmax()
        .item()
    )
    selected = [first]
    closest_sq = _pairwise_flat_sq(
        flattened,
        flattened[first:first + 1],
        norms_sq,
        norms_sq[first:first + 1],
    ).squeeze(1)
    for _ in range(1, num_landmarks):
        next_index = int(closest_sq.argmax().item())
        selected.append(next_index)
        next_sq = _pairwise_flat_sq(
            flattened,
            flattened[next_index:next_index + 1],
            norms_sq,
            norms_sq[next_index:next_index + 1],
        ).squeeze(1)
        closest_sq = torch.minimum(closest_sq, next_sq)

    centroids = samples.index_select(
        0,
        torch.tensor(selected, device=samples.device, dtype=torch.long),
    ).clone()
    previous_labels = None
    for _ in range(int(iterations)):
        labels, _ = _nearest_centroid_labels(
            samples, centroids, int(chunk_size))
        if previous_labels is not None and torch.equal(labels, previous_labels):
            break
        previous_labels = labels
        flattened = samples.flatten(1)
        sums, cluster_counts = _deterministic_cluster_sums(
            flattened, labels, num_landmarks)
        nonempty = cluster_counts > 0
        updated = centroids.flatten(1).clone()
        updated[nonempty] = (
            sums[nonempty]
            / cluster_counts[nonempty, None].to(sums.dtype)
        )
        centroids = updated.reshape(
            num_landmarks, num_tokens, feature_dim)
    return centroids


@torch.no_grad()
def compute_nystrom_alpha(
    reference: torch.Tensor,
    landmarks: torch.Tensor,
    sigma: torch.Tensor | float,
    *,
    ridge: float = 1.0e-6,
    chunk_size: int = 256,
) -> torch.Tensor:
    """Fit coefficients for the persistent Nyström reference kernel mean.

    Solve (K_ZZ + ridge * I) alpha = mean_R K_ZR once, offline. Training
    evaluates K(x, Z) @ alpha without revisiting the full reference bank R.
    Code ``alpha`` is the paper's coefficient vector beta (Appendix C.2),
    not the scalar mixture weight alpha_e in the hybrid mean (Eq. 9).
    """
    ref_count, ref_n, ref_d = _validate_video_tensor(reference, "reference")
    _, landmark_n, landmark_d = _validate_video_tensor(
        landmarks, "landmarks")
    if (ref_n, ref_d) != (landmark_n, landmark_d):
        raise ValueError("reference and landmarks have different matrix shapes")
    if not math.isfinite(float(ridge)) or float(ridge) < 0:
        raise ValueError("ridge must be finite and non-negative")
    if int(chunk_size) <= 0:
        raise ValueError("chunk_size must be positive")

    reference_flat = reference.detach().float().flatten(1)
    landmark_flat = landmarks.detach().float().flatten(1)
    reference_sq = _flat_row_norm_sq(reference_flat)
    landmark_sq = _flat_row_norm_sq(landmark_flat)
    inverse = _inverse_two_sigma_sq(sigma, device=landmarks.device)
    target = landmarks.new_zeros(landmarks.shape[0], dtype=torch.float64)
    for start in range(0, ref_count, int(chunk_size)):
        end = min(start + int(chunk_size), ref_count)
        block = _rbf_flat_precomputed(
            landmark_flat,
            reference_flat[start:end],
            landmark_sq,
            reference_sq[start:end],
            inverse,
        )
        target.add_(block.double().sum(dim=1))
    target.div_(float(ref_count))

    kernel_zz = _rbf_flat_precomputed(
        landmark_flat,
        landmark_flat,
        landmark_sq,
        landmark_sq,
        inverse,
    ).double()
    kernel_zz = 0.5 * (kernel_zz + kernel_zz.t())
    kernel_zz.fill_diagonal_(1.0)
    regularized = kernel_zz + float(ridge) * torch.eye(
        landmarks.shape[0],
        device=landmarks.device,
        dtype=torch.float64,
    )
    try:
        alpha = torch.linalg.solve(regularized, target)
    except RuntimeError as exc:
        raise ValueError(
            "Nyström landmark kernel is singular; increase ridge") from exc
    return alpha.float()


@torch.no_grad()
def compute_reference_krr(
    reference: torch.Tensor,
    sigma: torch.Tensor | float,
    *,
    chunk_size: int = 256,
) -> torch.Tensor:
    """Exact biased ``mean(K_rr)``, including the reference diagonal."""
    count, _, _ = _validate_video_tensor(reference, "reference")
    if int(chunk_size) <= 0:
        raise ValueError("chunk_size must be positive")
    flattened = reference.detach().float().flatten(1)
    norms_sq = _flat_row_norm_sq(flattened)
    inverse = _inverse_two_sigma_sq(sigma, device=reference.device)
    total = reference.new_zeros((), dtype=torch.float64)
    chunk_size = int(chunk_size)
    # The kernel is symmetric.  Evaluate the diagonal and strict upper block
    # triangle once, doubling only off-diagonal blocks.  This is the same full
    # biased Gram mean while avoiding millions of duplicate tiny GEMMs.
    for row_start in range(0, count, chunk_size):
        row_end = min(row_start + chunk_size, count)
        for col_start in range(row_start, count, chunk_size):
            col_end = min(col_start + chunk_size, count)
            kernel = _rbf_flat_precomputed(
                flattened[row_start:row_end],
                flattened[col_start:col_end],
                norms_sq[row_start:row_end],
                norms_sq[col_start:col_end],
                inverse,
            )
            block_sum = kernel.double().sum()
            if row_start == col_start:
                block_sum.add_(
                    (1.0 - kernel.diagonal()).double().sum())
                total.add_(block_sum)
            else:
                total.add_(2.0 * block_sum)
    return (total / float(count * count)).float()


def _hash_tensor(hasher: Any, name: str, tensor: torch.Tensor) -> None:
    value = tensor.detach().cpu().contiguous()
    descriptor = {
        "name": name,
        "dtype": str(value.dtype),
        "shape": list(value.shape),
    }
    hasher.update(json.dumps(
        descriptor, sort_keys=True, separators=(",", ":")).encode("utf-8"))
    raw = value.reshape(-1).view(torch.uint8).numpy().tobytes()
    hasher.update(raw)


def compute_reference_pack_hash(pack: Mapping[str, Any]) -> str:
    """Hash pack contents, deliberately excluding ``content_hash`` itself."""
    if not isinstance(pack, Mapping):
        raise TypeError("reference pack must be a mapping")
    required_without_hash = _PACK_KEYS - {"content_hash"}
    missing = required_without_hash - set(pack)
    if missing:
        raise ValueError(f"reference pack is missing keys: {sorted(missing)}")
    metadata = canonicalize_metadata(pack["metadata"])
    hasher = hashlib.sha256()
    hasher.update(b"video-nystrom-reference-pack\0")
    hasher.update(str(int(pack["format_version"])).encode("ascii"))
    hasher.update(json.dumps(
        metadata,
        sort_keys=True,
        separators=(",", ":"),
        ensure_ascii=True,
    ).encode("utf-8"))
    for name in ("landmarks", "alpha", "k_rr", "sigma"):
        value = pack[name]
        if not isinstance(value, torch.Tensor):
            raise TypeError(f"reference pack {name!r} must be a tensor")
        _hash_tensor(hasher, name, value)
    return hasher.hexdigest()


def validate_reference_pack(
    pack: Mapping[str, Any],
    *,
    expected_metadata: Optional[Mapping[str, Any]] = None,
) -> Dict[str, Any]:
    """Validate schema, geometry, recursive metadata subset, and SHA-256."""
    if not isinstance(pack, Mapping):
        raise TypeError("reference pack must be a mapping")
    keys = set(pack)
    if keys != _PACK_KEYS:
        raise ValueError(
            "reference pack keys mismatch: "
            f"missing={sorted(_PACK_KEYS - keys)}, "
            f"unexpected={sorted(keys - _PACK_KEYS)}")
    if int(pack["format_version"]) != REFERENCE_PACK_VERSION:
        raise ValueError(
            "unsupported reference pack version "
            f"{pack['format_version']!r}; expected {REFERENCE_PACK_VERSION}")

    landmarks = pack["landmarks"]
    alpha = pack["alpha"]
    k_rr = pack["k_rr"]
    sigma = pack["sigma"]
    landmark_count, num_tokens, feature_dim = _validate_video_tensor(
        landmarks, "landmarks")
    if not isinstance(alpha, torch.Tensor) or alpha.ndim != 1:
        raise ValueError("alpha must be a one-dimensional tensor")
    if alpha.shape[0] != landmark_count or not alpha.is_floating_point():
        raise ValueError("alpha length/dtype is incompatible with landmarks")
    if not bool(torch.isfinite(alpha).all().item()):
        raise FloatingPointError("alpha contains non-finite values")
    for name, value in (("k_rr", k_rr), ("sigma", sigma)):
        if (
            not isinstance(value, torch.Tensor)
            or value.numel() != 1
            or not value.is_floating_point()
        ):
            raise ValueError(f"{name} must be a floating scalar tensor")
        if not bool(torch.isfinite(value).item()):
            raise FloatingPointError(f"{name} is non-finite")
    if float(sigma) <= 0:
        raise ValueError("sigma must be positive")
    if float(k_rr) < 0:
        raise ValueError("k_rr must be non-negative")

    metadata = canonicalize_metadata(pack["metadata"])
    if metadata != pack["metadata"]:
        raise ValueError("reference pack metadata is not canonical")
    nystrom_metadata = metadata.get("nystrom")
    required_nystrom = {
        "sample_axis": "video",
        "sample_shape": [num_tokens, feature_dim],
        "kernel": "gaussian_rbf",
        "kernel_metric": "frobenius",
        "generated_gram": "biased_exact_including_diagonal",
        "num_landmarks": landmark_count,
    }
    if not isinstance(nystrom_metadata, Mapping):
        raise ValueError("metadata.nystrom is missing")
    _metadata_contains(
        nystrom_metadata, required_nystrom, "metadata.nystrom")
    if expected_metadata is not None:
        expected = canonicalize_metadata(expected_metadata)
        _metadata_contains(metadata, expected)

    saved_hash = pack["content_hash"]
    if not isinstance(saved_hash, str) or len(saved_hash) != 64:
        raise ValueError("content_hash must be a 64-character SHA-256 hex string")
    calculated_hash = compute_reference_pack_hash(pack)
    if not hmac.compare_digest(saved_hash, calculated_hash):
        raise ValueError(
            "reference pack content_hash mismatch: "
            f"saved={saved_hash}, calculated={calculated_hash}")

    return {
        "format_version": REFERENCE_PACK_VERSION,
        "landmarks": landmarks,
        "alpha": alpha,
        "k_rr": k_rr.reshape(()),
        "sigma": sigma.reshape(()),
        "metadata": metadata,
        "content_hash": saved_hash,
    }


@torch.no_grad()
def build_video_nystrom_reference(
    reference: torch.Tensor,
    *,
    num_landmarks: int,
    metadata: Mapping[str, Any],
    bandwidth: Optional[float] = None,
    median_max_samples: Optional[int] = 2048,
    kmeans_iterations: int = 20,
    ridge: float = 1.0e-6,
    chunk_size: int = 256,
    seed: int = 0,
) -> Dict[str, Any]:
    """Construct a frozen, portable Nyström reference pack."""
    reference_count, num_tokens, feature_dim = _validate_video_tensor(
        reference, "reference")
    if int(chunk_size) <= 0:
        raise ValueError("chunk_size must be positive")
    seed = int(seed)
    reference = reference.detach().float()
    if bandwidth is None:
        sigma = median_frobenius_sigma(
            reference,
            max_samples=median_max_samples,
            chunk_size=int(chunk_size),
        )
        bandwidth_source = "median_frobenius_distance"
    else:
        if not math.isfinite(float(bandwidth)) or float(bandwidth) <= 0:
            raise ValueError("bandwidth must be finite and positive")
        sigma = reference.new_tensor(float(bandwidth))
        bandwidth_source = "provided"

    landmarks = deterministic_video_kmeans(
        reference,
        int(num_landmarks),
        iterations=int(kmeans_iterations),
        chunk_size=int(chunk_size),
    )
    alpha = compute_nystrom_alpha(
        reference,
        landmarks,
        sigma,
        ridge=float(ridge),
        chunk_size=int(chunk_size),
    )
    k_rr = compute_reference_krr(
        reference, sigma, chunk_size=int(chunk_size))

    canonical_metadata = canonicalize_metadata(metadata)
    technical_metadata = {
        "sample_axis": "video",
        "sample_shape": [num_tokens, feature_dim],
        "kernel": "gaussian_rbf",
        "kernel_metric": "frobenius",
        "generated_gram": "biased_exact_including_diagonal",
        "reference_count": reference_count,
        "num_landmarks": int(num_landmarks),
        "bandwidth_source": bandwidth_source,
        "median_max_samples": (
            None if median_max_samples is None else int(median_max_samples)
        ),
        "kmeans": "deterministic_farthest_first_lloyd",
        "seed": seed,
        "kmeans_iterations": int(kmeans_iterations),
        "ridge": float(ridge),
    }
    if (
        "nystrom" in canonical_metadata
        and canonical_metadata["nystrom"] != technical_metadata
    ):
        raise ValueError(
            "metadata.nystrom is reserved and does not match builder settings")
    canonical_metadata["nystrom"] = technical_metadata
    canonical_metadata = canonicalize_metadata(canonical_metadata)

    pack: Dict[str, Any] = {
        "format_version": REFERENCE_PACK_VERSION,
        "landmarks": landmarks.detach().cpu().contiguous(),
        "alpha": alpha.detach().cpu().contiguous(),
        "k_rr": k_rr.detach().cpu().reshape(()),
        "sigma": sigma.detach().cpu().reshape(()),
        "metadata": canonical_metadata,
    }
    pack["content_hash"] = compute_reference_pack_hash(pack)
    return validate_reference_pack(pack)


def save_reference_pack(
    pack: Mapping[str, Any],
    path: str | os.PathLike,
    *,
    atomic: bool = True,
) -> None:
    """Validate and save a portable pack, atomically by default."""
    validated = validate_reference_pack(pack)
    portable = {
        **validated,
        "landmarks": validated["landmarks"].detach().cpu().contiguous(),
        "alpha": validated["alpha"].detach().cpu().contiguous(),
        "k_rr": validated["k_rr"].detach().cpu().reshape(()),
        "sigma": validated["sigma"].detach().cpu().reshape(()),
    }
    # Moving tensors to CPU is part of the serialized representation, so hash
    # the exact portable payload rather than copying the incoming hash blindly.
    portable["content_hash"] = compute_reference_pack_hash(portable)
    destination = os.path.abspath(os.fspath(path))
    if not atomic:
        torch.save(portable, destination)
        return
    directory = os.path.dirname(destination)
    descriptor, temporary = tempfile.mkstemp(
        prefix=f".{os.path.basename(destination)}.",
        suffix=".tmp",
        dir=directory,
    )
    os.close(descriptor)
    try:
        torch.save(portable, temporary)
        os.replace(temporary, destination)
    finally:
        if os.path.exists(temporary):
            os.unlink(temporary)


def load_reference_pack(
    path: str | os.PathLike,
    *,
    expected_metadata: Optional[Mapping[str, Any]] = None,
    map_location: str | torch.device = "cpu",
) -> Dict[str, Any]:
    """Load a pack and strictly validate its metadata and content hash."""
    try:
        pack = torch.load(
            os.fspath(path),
            map_location=map_location,
            weights_only=True,
        )
    except TypeError:
        pack = torch.load(os.fspath(path), map_location=map_location)
    return validate_reference_pack(
        pack, expected_metadata=expected_metadata)


def _kernel_sum_chunked(
    left: torch.Tensor,
    right: torch.Tensor,
    sigma: torch.Tensor,
    chunk_size: int,
) -> torch.Tensor:
    total = left.new_zeros(())
    self_kernel = left is right
    for row_start in range(0, left.shape[0], chunk_size):
        rows = left[row_start:row_start + chunk_size]
        for col_start in range(0, right.shape[0], chunk_size):
            cols = right[col_start:col_start + chunk_size]
            kernel = rbf_frobenius_kernel(rows, cols, sigma)
            block_sum = kernel.sum()
            if self_kernel and row_start == col_start:
                # The RBF self-kernel is analytically one.  Explicitly replace
                # it to avoid norm/Gram cancellation error in FP32.
                diagonal = kernel.diagonal()
                block_sum = block_sum + (1.0 - diagonal).sum()
            total = total + block_sum
    return total


def _kernel_matvec_chunked(
    samples: torch.Tensor,
    landmarks: torch.Tensor,
    alpha: torch.Tensor,
    sigma: torch.Tensor,
    chunk_size: int,
) -> torch.Tensor:
    outputs = []
    for row_start in range(0, samples.shape[0], chunk_size):
        rows = samples[row_start:row_start + chunk_size]
        values = rows.new_zeros(rows.shape[0])
        for col_start in range(0, landmarks.shape[0], chunk_size):
            cols = landmarks[col_start:col_start + chunk_size]
            kernel = rbf_frobenius_kernel(rows, cols, sigma)
            values = values + kernel @ alpha[
                col_start:col_start + chunk_size]
        outputs.append(values)
    return torch.cat(outputs)


class VideoNystromMMD(nn.Module):
    """One encoder's MMD: repulsion - 2 * attraction + reference constant.

    The generated-sample kernel is evaluated exactly. Nyström approximates
    only attraction to the persistent reference bank; attraction to the
    current reference minibatch is also exact when enabled.
    """

    _STATE_VERSION = 1
    _ATTRACTION_EMA_STATE_VERSION = 1
    _REPULSION_FIFO_STATE_VERSION = 1

    def __init__(
        self,
        reference_pack: Mapping[str, Any],
        *,
        expected_metadata: Optional[Mapping[str, Any]] = None,
        kernel_chunk_size: int = 256,
        attraction_ema_enabled: bool = False,
        attraction_ema_beta: float = 0.99,
        attraction_weight: float = 1.0,
        current_teacher_attraction_enabled: bool = False,
        current_teacher_attraction_weight: float = 1.0,
        repulsion_fifo_size: int = 0,
    ):
        super().__init__()
        if int(kernel_chunk_size) <= 0:
            raise ValueError("kernel_chunk_size must be positive")
        if (
            not math.isfinite(float(attraction_ema_beta))
            or not 0.0 <= float(attraction_ema_beta) < 1.0
        ):
            raise ValueError("attraction EMA beta must be in [0, 1)")
        if (
            not math.isfinite(float(attraction_weight))
            or float(attraction_weight) < 0.0
        ):
            raise ValueError("attraction weight must be finite and non-negative")
        if (
            not math.isfinite(float(current_teacher_attraction_weight))
            or float(current_teacher_attraction_weight) < 0.0
        ):
            raise ValueError(
                "current teacher attraction weight must be finite and "
                "non-negative")
        current_teacher_attraction_enabled = bool(
            current_teacher_attraction_enabled)
        if (
            current_teacher_attraction_enabled
            and float(attraction_weight)
            + float(current_teacher_attraction_weight) <= 0.0
        ):
            raise ValueError(
                "fixed/current teacher attraction weights must have a "
                "positive sum when current teacher attraction is enabled")
        repulsion_fifo_size = int(repulsion_fifo_size)
        if repulsion_fifo_size < 0:
            raise ValueError("repulsion FIFO size must be non-negative")
        if repulsion_fifo_size > 0 and attraction_ema_enabled:
            raise ValueError(
                "repulsion FIFO is incompatible with attraction EMA")
        pack = validate_reference_pack(
            reference_pack, expected_metadata=expected_metadata)
        self.kernel_chunk_size = int(kernel_chunk_size)
        self.attraction_ema_enabled = bool(attraction_ema_enabled)
        self.attraction_ema_beta = float(attraction_ema_beta)
        self.attraction_weight = float(attraction_weight)
        self.current_teacher_attraction_enabled = (
            current_teacher_attraction_enabled)
        self.current_teacher_attraction_weight = float(
            current_teacher_attraction_weight)
        self.repulsion_fifo_size = repulsion_fifo_size
        self.repulsion_fifo_enabled = bool(repulsion_fifo_size > 0)
        self.reference_metadata = pack["metadata"]
        self.reference_hash = pack["content_hash"]
        self.register_buffer(
            "landmarks", pack["landmarks"].detach().float().clone(),
            persistent=True)
        self.register_buffer(
            "alpha", pack["alpha"].detach().float().clone(),
            persistent=True)
        self.register_buffer(
            "k_rr", pack["k_rr"].detach().float().reshape(()).clone(),
            persistent=True)
        self.register_buffer(
            "sigma", pack["sigma"].detach().float().reshape(()).clone(),
            persistent=True)
        # FIFO contents have their own strict trainer-checkpoint payload.  A
        # non-persistent buffer still follows ``module.to(device)`` without
        # silently duplicating the large history in the module state_dict.
        self.register_buffer(
            "repulsion_fifo",
            torch.empty(
                (0, int(self.landmarks.shape[1]), int(self.landmarks.shape[2])),
                dtype=torch.float32,
                device=self.landmarks.device,
            ),
            persistent=False,
        )
        self.register_buffer(
            "repulsion_fifo_commits",
            torch.zeros((), dtype=torch.long, device=self.landmarks.device),
            persistent=False,
        )
        # Keep the disabled module's state_dict byte-schema compatible with
        # the original exact objective.  Enabled modules persist these two
        # scalars and advertise the changed schema through extra state.
        self.register_buffer(
            "attraction_ema",
            torch.zeros((), dtype=torch.float32),
            persistent=self.attraction_ema_enabled,
        )
        self.register_buffer(
            "attraction_ema_updates",
            torch.zeros((), dtype=torch.long),
            persistent=self.attraction_ema_enabled,
        )
        # Pending state is deliberately not a buffer: an uncommitted optimizer
        # update must never enter a checkpoint.
        self._pending_attraction: Optional[torch.Tensor] = None
        self._pending_repulsion_fifo: Optional[torch.Tensor] = None
        # The final FIFO allocation is prepared before optimizer.step().
        # Post-step commit only publishes this already-owned tensor.
        self._prepared_repulsion_fifo: Optional[torch.Tensor] = None

    @property
    def sample_shape(self) -> Tuple[int, int]:
        return int(self.landmarks.shape[1]), int(self.landmarks.shape[2])

    @property
    def attraction_ema_pending(self) -> bool:
        return self._pending_attraction is not None

    @property
    def repulsion_fifo_pending(self) -> bool:
        return self._pending_repulsion_fifo is not None

    @property
    def repulsion_fifo_fill(self) -> int:
        return int(self.repulsion_fifo.shape[0])

    def forward(
        self,
        generated: Optional[torch.Tensor] = None,
        current_teacher: Optional[torch.Tensor] = None,
        *,
        gradient_students: Optional[torch.Tensor] = None,
        student_reference: Optional[torch.Tensor] = None,
        gradient_student_indices: Optional[
            torch.Tensor | Sequence[int]
        ] = None,
        stage: bool = True,
    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
        """Evaluate the legacy full-batch or explicit partial-gradient loss.

        Legacy callers pass ``generated`` and retain the original two-sided
        ``K(generated, generated)`` derivative exactly. Partial-gradient
        callers instead pass both ``gradient_students`` (the row-side samples
        whose derivatives are required) and ``student_reference`` (the full
        detached column-side student population). The latter computes

        ``mean K(G,S) - 2 * attraction(G) + reference_constant``

        with a rectangular ``G x S`` student population. The matched
        teacher population must have exactly ``|G|`` samples. Neither
        ``student_reference`` nor ``current_teacher`` can receive gradients.

        Passing ``gradient_student_indices`` opts into the full-student
        partial-gradient estimator. The indices identify the positions of
        ``gradient_students`` in ``student_reference``. Those differentiable
        values are functionally inserted into the otherwise detached student
        population, yielding complete ``mean K(S,S)``, fixed attraction over
        all ``S``, and exact current-teacher attraction over ``S x T``. Only
        indexed students retain gradients; ``|T|`` must still equal ``|G|``.
        """
        # 1. Choose the generated population and samples that retain a graph.
        # ``student_reference`` is generated data, not the persistent reference
        # bank; these input names remain compatible with existing callers.
        partial_mode = (
            gradient_students is not None or student_reference is not None)
        full_student_partial_mode = gradient_student_indices is not None
        if full_student_partial_mode and not partial_mode:
            raise ValueError(
                "gradient_student_indices is only valid with the explicit "
                "partial-gradient interface")
        if generated is not None and gradient_students is not None:
            raise ValueError(
                "generated and gradient_students are mutually exclusive")
        if partial_mode:
            if gradient_students is None or student_reference is None:
                raise ValueError(
                    "partial-gradient MMD requires both gradient_students and "
                    "student_reference")
            if generated is not None:
                raise ValueError(
                    "generated cannot be combined with the explicit "
                    "partial-gradient interface")
            if self.repulsion_fifo_enabled:
                raise ValueError(
                    "partial-gradient student_reference is incompatible with "
                    "repulsion FIFO")
            if self.attraction_ema_enabled:
                raise ValueError(
                    "partial-gradient student_reference is incompatible with "
                    "attraction EMA")
            gradient_input = gradient_students
        else:
            if generated is None:
                raise ValueError(
                    "generated is required for legacy MMD calls")
            gradient_input = generated

        gradient_batch, num_tokens, feature_dim = _validate_video_tensor(
            gradient_input,
            "gradient_students" if partial_mode else "generated",
        )
        if (num_tokens, feature_dim) != self.sample_shape:
            raise ValueError(
                "student video matrix shape does not match frozen reference: "
                f"{(num_tokens, feature_dim)} != {self.sample_shape}")
        if gradient_input.device != self.landmarks.device:
            raise ValueError(
                "gradient students and frozen reference must be on the same "
                "device")

        samples = gradient_input.float()
        selected_indices = None
        if partial_mode:
            reference_batch, reference_tokens, reference_dim = (
                _validate_video_tensor(
                    student_reference, "student_reference"))
            if (reference_tokens, reference_dim) != (num_tokens, feature_dim):
                raise ValueError(
                    "student_reference matrix shape must match "
                    "gradient_students: "
                    f"{(reference_tokens, reference_dim)} != "
                    f"{(num_tokens, feature_dim)}")
            if student_reference.device != gradient_input.device:
                raise ValueError(
                    "gradient_students and student_reference must share a "
                    "device")
            reference_samples = student_reference.detach().float()
            if full_student_partial_mode:
                try:
                    selected_indices = torch.as_tensor(
                        gradient_student_indices,
                        device=gradient_input.device,
                    )
                except (TypeError, ValueError) as exc:
                    raise TypeError(
                        "gradient_student_indices must be an integer "
                        "tensor or sequence") from exc
                integer_dtypes = (
                    torch.uint8,
                    torch.int8,
                    torch.int16,
                    torch.int32,
                    torch.int64,
                )
                if (
                    selected_indices.dtype == torch.bool
                    or selected_indices.dtype not in integer_dtypes
                ):
                    raise TypeError(
                        "gradient_student_indices must contain integers")
                selected_indices = selected_indices.to(dtype=torch.long)
                if selected_indices.ndim != 1:
                    raise ValueError(
                        "gradient_student_indices must be one-dimensional")
                if int(selected_indices.numel()) != gradient_batch:
                    raise ValueError(
                        "gradient_student_indices must contain exactly one "
                        "index per gradient student: "
                        f"{int(selected_indices.numel())} != {gradient_batch}")
                if bool(
                    ((selected_indices < 0)
                     | (selected_indices >= reference_batch)).any().item()
                ):
                    raise ValueError(
                        "gradient_student_indices contains an out-of-range "
                        "student_reference position")
                if (
                    int(torch.unique(selected_indices).numel())
                    != gradient_batch
                ):
                    raise ValueError(
                        "gradient_student_indices must be unique")
                selected_reference = reference_samples.index_select(
                    0, selected_indices)
                if not torch.equal(samples.detach(), selected_reference):
                    raise ValueError(
                        "gradient_students must exactly match "
                        "student_reference at gradient_student_indices")
                # Functionally replace only selected detached values. Reusing
                # this same object on both sides of K(S,S) gives selected
                # students both row- and column-side derivatives.
                objective_samples = reference_samples.index_copy(
                    0, selected_indices, samples)
                repulsion_reference_samples = objective_samples
                attraction_samples = objective_samples
                attraction_batch = reference_batch
            else:
                objective_samples = samples
                repulsion_reference_samples = reference_samples
                attraction_samples = samples
                attraction_batch = gradient_batch
        else:
            reference_batch = gradient_batch
            # Preserve object identity here. The legacy self-kernel helper
            # relies on it both for its exact diagonal replacement and for the
            # original two-sided K_ss derivative.
            reference_samples = samples
            objective_samples = samples
            repulsion_reference_samples = reference_samples
            attraction_samples = samples
            attraction_batch = gradient_batch

        # 2. Read the frozen reference minibatch for the Monte Carlo term.
        # These features are targets; gradients flow through generated samples.
        teacher_samples = None
        teacher_batch = None
        if self.current_teacher_attraction_enabled:
            if current_teacher is None:
                raise ValueError(
                    "current_teacher is required when exact current-teacher "
                    "attraction is enabled")
            teacher_batch, teacher_tokens, teacher_dim = (
                _validate_video_tensor(current_teacher, "current_teacher"))
            if (
                teacher_batch != gradient_batch
                or teacher_tokens != num_tokens
                or teacher_dim != feature_dim
            ):
                student_name = (
                    "gradient_students" if partial_mode else "generated")
                raise ValueError(
                    "current teacher features must exactly match "
                    f"{student_name} "
                    "[B,N,D]: "
                    f"{(teacher_batch, teacher_tokens, teacher_dim)} != "
                    f"{(gradient_batch, num_tokens, feature_dim)}")
            if current_teacher.device != gradient_input.device:
                raise ValueError(
                    "gradient students and current teacher must be on the same "
                    "device")
            # The current teacher is a matched frozen target. Detach before
            # every kernel operation so callers can never accidentally create
            # a teacher-side gradient graph.
            teacher_samples = current_teacher.detach().float()
        elif current_teacher is not None:
            raise ValueError(
                "current_teacher was provided while current teacher "
                "attraction is disabled")

        # 3. Exact generated-sample repulsion: sum every kernel pair in
        # chunks. Chunking controls memory without approximating the sum.
        exact_kgg_sum = _kernel_sum_chunked(
            objective_samples,
            repulsion_reference_samples,
            self.sigma,
            self.kernel_chunk_size,
        )
        if self.repulsion_fifo_enabled:
            history_size = self.repulsion_fifo_fill
            denominator = gradient_batch * (
                gradient_batch - 1 + history_size)
            if denominator <= 0:
                raise ValueError(
                    "repulsion FIFO needs at least one current-current "
                    "off-diagonal or current-history pair")
            # ``_kernel_sum_chunked`` analytically replaces each self-kernel
            # by one.  Subtracting B therefore removes the complete diagonal,
            # including its autograd path, without forming a dense mask.
            current_offdiag_sum = exact_kgg_sum - float(gradient_batch)
            if history_size:
                history_cross_sum = _kernel_sum_chunked(
                    samples,
                    self.repulsion_fifo,
                    self.sigma,
                    self.kernel_chunk_size,
                )
            else:
                history_cross_sum = samples.new_zeros(())
            repulsion = (
                current_offdiag_sum + history_cross_sum
            ) / float(denominator)
        else:
            history_size = 0
            # Full-population replay keeps B * B pairs even when only a
            # subset is differentiated. The rectangular mode uses |G| * B.
            # Gradient-subset compensation is handled by the replay caller.
            if full_student_partial_mode:
                denominator = reference_batch * reference_batch
            else:
                denominator = gradient_batch * reference_batch
            current_offdiag_sum = None
            history_cross_sum = None
            repulsion = exact_kgg_sum / float(denominator)
        # 4. Persistent Nyström attraction (paper Eq. 8): K(x, Z) @ alpha.
        # The saved ``alpha`` vector is paper beta, not the mixture weight.
        fixed_attraction_values = _kernel_matvec_chunked(
            attraction_samples,
            self.landmarks,
            self.alpha,
            self.sigma,
            self.kernel_chunk_size,
        )
        fixed_current_attraction = fixed_attraction_values.mean()
        current_teacher_attraction = None
        current_teacher_self_kernel = None
        current_teacher_exact_mmd = None
        # 5. Hybrid reference mean (paper Eq. 9): mix the persistent term
        # with exact minibatch attraction over all generated x reference pairs.
        # Group weights 2:1 correspond to paper alpha_e = 1/3.
        if self.current_teacher_attraction_enabled:
            fixed_weight = self.attraction_weight
            current_weight = self.current_teacher_attraction_weight
            teacher_weight_sum = fixed_weight + current_weight
            fixed_normalized_weight = fixed_weight / teacher_weight_sum
            current_normalized_weight = current_weight / teacher_weight_sum
            current_teacher_attraction = _kernel_sum_chunked(
                attraction_samples,
                teacher_samples,
                self.sigma,
                self.kernel_chunk_size,
            ) / float(attraction_batch * teacher_batch)
            # The reference self-kernel has no generator gradient. Its value
            # still affects the release's later detached-scale division.
            with torch.no_grad():
                current_teacher_self_kernel = _kernel_sum_chunked(
                    teacher_samples,
                    teacher_samples,
                    self.sigma,
                    self.kernel_chunk_size,
                ) / float(teacher_batch * teacher_batch)
            current_attraction = (
                fixed_normalized_weight * fixed_current_attraction
                + current_normalized_weight * current_teacher_attraction
            )
            reference_constant = (
                fixed_normalized_weight * self.k_rr
                + current_normalized_weight * current_teacher_self_kernel
            )
            current_teacher_exact_mmd = (
                repulsion
                - 2.0 * current_teacher_attraction
                + current_teacher_self_kernel
            )
        else:
            fixed_normalized_weight = None
            current_normalized_weight = None
            teacher_weight_sum = None
            current_attraction = fixed_current_attraction
            reference_constant = self.k_rr
        # 6. Optional historical attraction/FIFO state. These branches are
        # inactive in the released multi-encoder configuration.
        history_updates = int(self.attraction_ema_updates.item())
        if self.attraction_ema_enabled and history_updates > 0:
            ema_blend_current_coefficient = 1.0 - self.attraction_ema_beta
            ema_attraction = (
                self.attraction_ema_beta * self.attraction_ema.detach()
                + ema_blend_current_coefficient * current_attraction
            )
            if self.current_teacher_attraction_enabled:
                # Keep the transactional EMA's smoothed forward value while
                # preserving the full derivative of this update's normalized
                # fixed/current teacher mixture. Without this straight-through
                # form beta=0.99 would silently shrink all teacher gradients
                # to one percent relative to K_ss.
                attraction = current_attraction + (
                    ema_attraction - current_attraction).detach()
                current_coefficient = 1.0
            else:
                attraction = ema_attraction
                current_coefficient = ema_blend_current_coefficient
        else:
            ema_blend_current_coefficient = 1.0
            current_coefficient = 1.0
            attraction = current_attraction

        if self.attraction_ema_enabled and stage:
            if self._pending_attraction is not None:
                raise RuntimeError(
                    "an attraction EMA update is already staged; "
                    "commit or discard it before another staged forward")
            self._pending_attraction = current_attraction.detach().clone()

        if self.repulsion_fifo_enabled and stage:
            if self._pending_repulsion_fifo is not None:
                raise RuntimeError(
                    "a repulsion FIFO update is already staged; "
                    "commit or discard it before another staged forward")
            # The fresh GradCache supplies the identical global batch on every
            # rank.  Cloning it here gives each rank the same independently
            # owned FIFO candidate and prevents all gradients/aliasing.
            self._pending_repulsion_fifo = (
                samples.detach().to(dtype=torch.float32).clone())

        # 7. Assemble the raw single-encoder loss. MultiEncoderMMD applies
        # detached-scale normalization and encoder weighting after this.
        if self.current_teacher_attraction_enabled:
            # This is the weighted average of the fixed-reference MMD and the
            # current exact MMD. K_ss is shared, so it appears once and the
            # total attraction coefficient remains -2 for every raw scale.
            loss = repulsion - 2.0 * attraction + reference_constant
            effective_attraction_weight = 1.0
        else:
            # Preserve the legacy value and gradient exactly when the opt-in
            # current teacher path is disabled.
            if self.attraction_weight == 1.0:
                loss = repulsion - 2.0 * attraction + self.k_rr
            else:
                loss = (
                    repulsion
                    - 2.0 * self.attraction_weight * attraction
                    + self.k_rr
                )
            effective_attraction_weight = self.attraction_weight
        metadata = {
            "repulsion": repulsion.detach(),
            "attraction": attraction.detach(),
            "attraction_weight": loss.new_tensor(
                effective_attraction_weight),
            "reference_constant": reference_constant.detach(),
            "sigma": self.sigma.detach(),
            "global_samples": torch.tensor(
                reference_batch, device=loss.device, dtype=torch.long),
        }
        if partial_mode:
            metadata.update({
                "partial_student_reference_enabled": loss.new_ones(
                    (), dtype=torch.long),
                "gradient_student_samples": loss.new_tensor(
                    gradient_batch, dtype=torch.long),
                "student_reference_samples": loss.new_tensor(
                    reference_batch, dtype=torch.long),
                "student_reference_detached": loss.new_ones(
                    (), dtype=torch.long),
                "repulsion_denominator": loss.new_tensor(
                    denominator, dtype=torch.long),
                "student_repulsion_denominator": loss.new_tensor(
                    denominator, dtype=torch.long),
                "fixed_teacher_attraction_denominator": loss.new_tensor(
                    attraction_batch, dtype=torch.long),
            })
            if full_student_partial_mode:
                metadata.update({
                    "full_student_partial_gradient": loss.new_ones(
                        (), dtype=torch.long),
                    "student_repulsion_gradient_sides": loss.new_tensor(
                        2, dtype=torch.long),
                })
        if self.current_teacher_attraction_enabled:
            metadata.update({
                "current_teacher_attraction_enabled": loss.new_ones(
                    (), dtype=torch.long),
                "fixed_teacher_attraction": (
                    fixed_current_attraction.detach()),
                "current_teacher_attraction": (
                    current_teacher_attraction.detach()),
                "fixed_teacher_attraction_weight": loss.new_tensor(
                    self.attraction_weight),
                "current_teacher_attraction_weight": loss.new_tensor(
                    self.current_teacher_attraction_weight),
                "teacher_attraction_weight_sum": loss.new_tensor(
                    teacher_weight_sum),
                "fixed_teacher_attraction_normalized_weight": (
                    loss.new_tensor(fixed_normalized_weight)),
                "current_teacher_attraction_normalized_weight": (
                    loss.new_tensor(current_normalized_weight)),
                "mixed_current_attraction": current_attraction.detach(),
                "fixed_reference_constant": self.k_rr.detach(),
                "current_teacher_self_kernel": (
                    current_teacher_self_kernel.detach()),
                "current_teacher_exact_mmd": (
                    current_teacher_exact_mmd.detach()),
                "current_teacher_samples": loss.new_tensor(
                    teacher_batch, dtype=torch.long),
            })
            if partial_mode:
                metadata.update({
                    "current_teacher_attraction_denominator": (
                        loss.new_tensor(
                            attraction_batch * teacher_batch,
                            dtype=torch.long,
                        )
                    ),
                    "current_teacher_self_kernel_denominator": (
                        loss.new_tensor(
                            teacher_batch * teacher_batch,
                            dtype=torch.long,
                        )
                    ),
                })
        if self.attraction_ema_enabled:
            metadata.update({
                "current_attraction": current_attraction.detach(),
                "attraction_ema_enabled": loss.new_ones(
                    (), dtype=torch.long),
                "attraction_ema_updates": (
                    self.attraction_ema_updates.detach().clone()),
                "attraction_current_coefficient": loss.new_tensor(
                    current_coefficient),
            })
            if self.current_teacher_attraction_enabled:
                metadata.update({
                    "attraction_ema_blend_current_coefficient": (
                        loss.new_tensor(
                            ema_blend_current_coefficient)),
                    "attraction_gradient_current_coefficient": (
                        loss.new_ones(())),
                    "attraction_ema_smoothed_value": attraction.detach(),
                    "effective_ema_loss": loss.detach(),
                    "attraction_gradient_mode": (
                        "straight_through_full_current_mix"),
                })
        if self.repulsion_fifo_enabled:
            fill_after_commit = min(
                self.repulsion_fifo_size,
                history_size + gradient_batch,
            )
            metadata.update({
                "repulsion_fifo_enabled": loss.new_ones(
                    (), dtype=torch.long),
                "repulsion_fifo_capacity": loss.new_tensor(
                    self.repulsion_fifo_size, dtype=torch.long),
                "repulsion_fifo_fill_before": loss.new_tensor(
                    history_size, dtype=torch.long),
                # Compatibility spelling used by the trainer logger.
                "repulsion_fifo_filled_before": loss.new_tensor(
                    history_size, dtype=torch.long),
                "repulsion_fifo_fill_after_commit": loss.new_tensor(
                    fill_after_commit, dtype=torch.long),
                "repulsion_fifo_commits": (
                    self.repulsion_fifo_commits.detach().clone()),
                "repulsion_fifo_committed_updates": (
                    self.repulsion_fifo_commits.detach().clone()),
                "repulsion_current_offdiag_kernel_sum": (
                    current_offdiag_sum.detach()),
                "repulsion_fifo_current_offdiag_sum": (
                    current_offdiag_sum.detach()),
                "repulsion_history_kernel_sum": (
                    history_cross_sum.detach()),
                "repulsion_fifo_history_cross_sum": (
                    history_cross_sum.detach()),
                "repulsion_denominator": loss.new_tensor(
                    denominator, dtype=torch.long),
                "repulsion_fifo_denominator": loss.new_tensor(
                    denominator, dtype=torch.long),
            })
        return loss, metadata

    @torch.no_grad()
    def prepare_commit(self) -> None:
        """Validate all staged distribution state before optimizer mutation."""
        if self.attraction_ema_enabled:
            if self._pending_attraction is None:
                raise RuntimeError("no attraction EMA update is staged")
            if (
                self._pending_attraction.ndim != 0
                or not bool(torch.isfinite(self._pending_attraction).item())
            ):
                raise FloatingPointError(
                    "refusing to commit a non-finite attraction EMA update")
        if self.repulsion_fifo_enabled:
            pending = self._pending_repulsion_fifo
            if pending is None:
                raise RuntimeError("no repulsion FIFO update is staged")
            if (
                pending.ndim != 3
                or tuple(pending.shape[1:]) != self.sample_shape
                or pending.shape[0] <= 0
                or pending.dtype != torch.float32
                or pending.device != self.repulsion_fifo.device
                or pending.requires_grad
                or pending.grad_fn is not None
            ):
                raise ValueError(
                    "staged repulsion FIFO values must be detached FP32 "
                    f"[B,{self.sample_shape[0]},{self.sample_shape[1]}] "
                    "on the FIFO device")
            if not bool(torch.isfinite(pending).all().item()):
                raise FloatingPointError(
                    "refusing to commit non-finite repulsion FIFO values")
            if self._prepared_repulsion_fifo is None:
                combined = torch.cat((self.repulsion_fifo, pending), dim=0)
                if combined.shape[0] > self.repulsion_fifo_size:
                    combined = combined[-self.repulsion_fifo_size:]
                # Allocate the right-sized, independently owned result while
                # the trainer can still abort before optimizer mutation.
                self._prepared_repulsion_fifo = (
                    combined.detach().contiguous().clone())

    @torch.no_grad()
    def commit(self) -> None:
        """Publish staged state only after a successful optimizer update."""
        if (
            self._pending_attraction is None
            and self._pending_repulsion_fifo is None
        ):
            return
        self.prepare_commit()
        if self._pending_attraction is not None:
            pending_attraction = self._pending_attraction
            if int(self.attraction_ema_updates.item()) == 0:
                self.attraction_ema.copy_(pending_attraction)
            else:
                self.attraction_ema.mul_(self.attraction_ema_beta).add_(
                    pending_attraction,
                    alpha=1.0 - self.attraction_ema_beta,
                )
            self.attraction_ema_updates.add_(1)
        if self._pending_repulsion_fifo is not None:
            if self._prepared_repulsion_fifo is None:
                raise RuntimeError("repulsion FIFO commit was not prepared")
            self.repulsion_fifo = self._prepared_repulsion_fifo
            self.repulsion_fifo_commits.add_(1)
        self.discard()

    def discard(self) -> None:
        """Drop all staged state after a failed or skipped update."""
        self._pending_attraction = None
        self._pending_repulsion_fifo = None
        self._prepared_repulsion_fifo = None

    def repulsion_fifo_state_dict(self) -> Optional[Dict[str, Any]]:
        """Return the committed FIFO in oldest-to-newest order."""
        if not self.repulsion_fifo_enabled:
            return None
        if self.repulsion_fifo_pending:
            raise RuntimeError(
                "refusing to checkpoint a pending repulsion FIFO update")
        return {
            "version": self._REPULSION_FIFO_STATE_VERSION,
            "capacity": self.repulsion_fifo_size,
            "sample_shape": list(self.sample_shape),
            "dtype": "float32",
            "count": self.repulsion_fifo_fill,
            "commits": int(self.repulsion_fifo_commits.item()),
            "values": self.repulsion_fifo.detach().cpu().clone(),
        }
    def validate_repulsion_fifo_state_dict(
        self,
        state: Mapping[str, Any],
    ) -> Dict[str, Any]:
        """Validate and normalize a portable FIFO checkpoint payload."""
        if not self.repulsion_fifo_enabled:
            raise ValueError(
                "cannot load repulsion FIFO state while it is disabled")
        if not isinstance(state, Mapping):
            raise TypeError("repulsion FIFO state must be a mapping")
        expected_keys = {
            "version", "capacity", "sample_shape", "dtype",
            "count", "commits", "values",
        }
        if set(state) != expected_keys:
            raise ValueError(
                "repulsion FIFO state keys mismatch: "
                f"saved={sorted(state)}, expected={sorted(expected_keys)}")
        if (
            state["version"] != self._REPULSION_FIFO_STATE_VERSION
            or int(state["capacity"]) != self.repulsion_fifo_size
            or tuple(state["sample_shape"]) != self.sample_shape
            or state["dtype"] != "float32"
        ):
            raise ValueError("repulsion FIFO checkpoint/config mismatch")
        count = int(state["count"])
        commits = int(state["commits"])
        values = state["values"]
        if not 0 <= count <= self.repulsion_fifo_size or commits < 0:
            raise ValueError("repulsion FIFO counters are invalid")
        if (
            not isinstance(values, torch.Tensor)
            or values.dtype != torch.float32
            or values.layout != torch.strided
            or values.is_meta
            or tuple(values.shape) != (count, *self.sample_shape)
            or not bool(torch.isfinite(values).all().item())
        ):
            raise ValueError(
                "repulsion FIFO values must be finite float32 with shape "
                f"{(count, *self.sample_shape)}")
        return {
            "version": self._REPULSION_FIFO_STATE_VERSION,
            "capacity": self.repulsion_fifo_size,
            "sample_shape": list(self.sample_shape),
            "dtype": "float32",
            "count": count,
            "commits": commits,
            "values": values.detach().cpu().clone(),
        }

    @torch.no_grad()
    def load_repulsion_fifo_state_dict(
        self,
        state: Mapping[str, Any],
    ) -> None:
        validated = self.validate_repulsion_fifo_state_dict(state)
        self.repulsion_fifo = validated["values"].to(
            device=self.landmarks.device,
            dtype=torch.float32,
        ).contiguous().clone()
        self.repulsion_fifo_commits.fill_(validated["commits"])
        self.discard()

    def attraction_ema_state_dict(self) -> Optional[Dict[str, Any]]:
        """Return the small trainer checkpoint payload for the opt-in EMA."""
        if not self.attraction_ema_enabled:
            return None
        if self.attraction_ema_pending:
            raise RuntimeError(
                "refusing to checkpoint a pending attraction EMA update")
        return {
            "version": self._ATTRACTION_EMA_STATE_VERSION,
            "beta": self.attraction_ema_beta,
            "value": self.attraction_ema.detach().cpu().clone(),
            "updates": self.attraction_ema_updates.detach().cpu().clone(),
        }

    def validate_attraction_ema_state_dict(
        self,
        state: Mapping[str, Any],
    ) -> Dict[str, Any]:
        """Validate and normalize a small attraction EMA checkpoint payload."""
        if not self.attraction_ema_enabled:
            raise ValueError(
                "cannot load attraction EMA state while the path is disabled")
        if not isinstance(state, Mapping):
            raise TypeError("attraction EMA state must be a mapping")
        expected_keys = {"version", "beta", "value", "updates"}
        if set(state) != expected_keys:
            raise ValueError(
                "attraction EMA state keys mismatch: "
                f"saved={sorted(state)}, expected={sorted(expected_keys)}")
        if (
            state["version"] != self._ATTRACTION_EMA_STATE_VERSION
            or float(state["beta"]) != self.attraction_ema_beta
        ):
            raise ValueError(
                "attraction EMA checkpoint/config mismatch: "
                f"saved version/beta={state['version']!r}/{state['beta']!r}, "
                "expected "
                f"{self._ATTRACTION_EMA_STATE_VERSION}/"
                f"{self.attraction_ema_beta}")
        value = state["value"]
        updates = state["updates"]
        if (
            not isinstance(value, torch.Tensor)
            or value.ndim != 0
            or value.dtype != torch.float32
            or value.layout != torch.strided
            or value.is_meta
            or not bool(torch.isfinite(value).item())
        ):
            raise ValueError(
                "attraction EMA checkpoint value must be a finite scalar "
                "torch.float32 tensor")
        if (
            not isinstance(updates, torch.Tensor)
            or updates.ndim != 0
            or updates.dtype != torch.long
            or updates.layout != torch.strided
            or updates.is_meta
            or int(updates.item()) < 0
        ):
            raise ValueError(
                "attraction EMA checkpoint updates must be a non-negative "
                "scalar torch.long tensor")
        if int(updates.item()) == 0 and float(value.item()) != 0.0:
            raise ValueError(
                "an uninitialized attraction EMA must have value zero")
        return {
            "version": self._ATTRACTION_EMA_STATE_VERSION,
            "beta": self.attraction_ema_beta,
            "value": value.detach().cpu().clone(),
            "updates": updates.detach().cpu().clone(),
        }

    @torch.no_grad()
    def load_attraction_ema_state_dict(
        self,
        state: Mapping[str, Any],
    ) -> None:
        """Strictly restore only the transactional attraction EMA scalars."""
        validated = self.validate_attraction_ema_state_dict(state)
        self.attraction_ema.copy_(
            validated["value"].to(
                device=self.attraction_ema.device,
                dtype=self.attraction_ema.dtype,
            )
        )
        self.attraction_ema_updates.copy_(
            validated["updates"].to(
                device=self.attraction_ema_updates.device)
        )
        self.discard()

    def get_extra_state(self) -> Dict[str, Any]:
        state = {
            "version": self._STATE_VERSION,
            "reference_metadata": self.reference_metadata,
            "reference_hash": self.reference_hash,
        }
        if self.attraction_ema_enabled:
            state = {
                **state,
                "version": self._STATE_VERSION + 1,
                "attraction_ema": {
                    "version": self._ATTRACTION_EMA_STATE_VERSION,
                    "beta": self.attraction_ema_beta,
                },
            }
        if self.current_teacher_attraction_enabled:
            state = {
                **state,
                "version": (4 if self.attraction_ema_enabled else 3),
                "current_teacher_attraction": {
                    "enabled": True,
                    "fixed_weight": self.attraction_weight,
                    "current_weight": (
                        self.current_teacher_attraction_weight),
                    "normalization": "divide_by_group_weight_sum",
                    "estimator": "exact_biased_gaussian_mmd",
                },
            }
        return state

    def set_extra_state(self, state: Mapping[str, Any]) -> None:
        expected = self.get_extra_state()
        if state != expected:
            raise ValueError(
                "Video Nyström MMD checkpoint/reference mismatch: "
                f"saved={state}, current={expected}")