from __future__ import annotations

import argparse
import json
import math
import random
from copy import deepcopy
from dataclasses import dataclass
from typing import Callable

import torch
from torch import Tensor, nn
from torch.nn import functional as F


@dataclass(frozen=True)
class ToyConfig:
    image_size: int = 32
    patch_size: int = 4
    channels: int = 3
    width: int = 64
    depth: int = 2
    heads: int = 4
    mlp_ratio: int = 4
    classes: int = 10
    batch_size: int = 2

    def __post_init__(self) -> None:
        for name, value in self.__dict__.items():
            assert isinstance(value, int) and value > 0, f"{name} must be positive"
        assert self.image_size % self.patch_size == 0
        assert self.width % self.heads == 0

    @property
    def grid_size(self) -> int:
        return self.image_size // self.patch_size

    @property
    def patch_tokens(self) -> int:
        return self.grid_size**2


def seed_everything(seed: int) -> None:
    random.seed(seed)
    torch.manual_seed(seed)
    torch.set_num_threads(1)
    torch.use_deterministic_algorithms(True)


class PatchEmbedding(nn.Module):
    def __init__(self, config: ToyConfig) -> None:
        super().__init__()
        self.config = config
        self.proj = nn.Conv2d(
            config.channels,
            config.width,
            kernel_size=config.patch_size,
            stride=config.patch_size,
        )

    def forward(self, images: Tensor) -> Tensor:
        assert images.ndim == 4
        batch, channels, height, width = images.shape
        assert channels == self.config.channels
        assert height == self.config.image_size
        assert width == self.config.image_size
        patches = self.proj(images).flatten(2).transpose(1, 2)
        assert patches.shape == (
            batch,
            self.config.patch_tokens,
            self.config.width,
        )
        return patches


class MultiHeadSelfAttention(nn.Module):
    def __init__(self, width: int, heads: int) -> None:
        super().__init__()
        assert width % heads == 0
        self.heads = heads
        self.head_dim = width // heads
        self.scale = self.head_dim**-0.5
        self.qkv = nn.Linear(width, 3 * width)
        self.proj = nn.Linear(width, width)

    def forward(
        self, tokens: Tensor, key_sizes: Tensor | None = None
    ) -> Tensor:
        assert tokens.ndim == 3
        batch, count, width = tokens.shape
        qkv = self.qkv(tokens)
        qkv = qkv.reshape(batch, count, 3, self.heads, self.head_dim)
        qkv = qkv.permute(2, 0, 3, 1, 4)
        query, key, value = qkv.unbind(0)
        scores = query @ key.transpose(-2, -1) * self.scale
        assert scores.shape == (batch, self.heads, count, count)
        if key_sizes is not None:
            assert key_sizes.shape == (batch, count)
            assert key_sizes.device == tokens.device
            assert key_sizes.dtype == tokens.dtype
            assert bool(torch.isfinite(key_sizes).all().item())
            assert bool((key_sizes > 0).all().item())
            scores = scores + key_sizes.log()[:, None, None, :]
        probabilities = scores.softmax(dim=-1)
        mixed = probabilities @ value
        mixed = mixed.transpose(1, 2).reshape(batch, count, width)
        output = self.proj(mixed)
        assert output.shape == tokens.shape
        return output


class MLP(nn.Module):
    def __init__(self, width: int, mlp_ratio: int) -> None:
        super().__init__()
        hidden = width * mlp_ratio
        self.fc1 = nn.Linear(width, hidden)
        self.act = nn.GELU()
        self.fc2 = nn.Linear(hidden, width)

    def forward(self, tokens: Tensor) -> Tensor:
        output = self.fc2(self.act(self.fc1(tokens)))
        assert output.shape == tokens.shape
        return output


class EncoderBlock(nn.Module):
    def __init__(self, width: int, heads: int, mlp_ratio: int) -> None:
        super().__init__()
        self.norm1 = nn.LayerNorm(width)
        self.attn = MultiHeadSelfAttention(width, heads)
        self.norm2 = nn.LayerNorm(width)
        self.mlp = MLP(width, mlp_ratio)

    def forward(
        self, tokens: Tensor, key_sizes: Tensor | None = None
    ) -> Tensor:
        tokens = tokens + self.attn(self.norm1(tokens), key_sizes)
        tokens = tokens + self.mlp(self.norm2(tokens))
        return tokens


class TinyViT(nn.Module):
    def __init__(self, config: ToyConfig, extra_tokens: int = 1) -> None:
        super().__init__()
        assert extra_tokens >= 1
        self.config = config
        self.extra_tokens = extra_tokens
        self.patch_embed = PatchEmbedding(config)
        self.cls_token = nn.Parameter(torch.zeros(1, 1, config.width))
        self.position = nn.Parameter(
            torch.zeros(1, config.patch_tokens + extra_tokens, config.width)
        )
        self.blocks = nn.ModuleList(
            [
                EncoderBlock(config.width, config.heads, config.mlp_ratio)
                for _ in range(config.depth)
            ]
        )
        self.norm = nn.LayerNorm(config.width)
        self.head = nn.Linear(config.width, config.classes)
        nn.init.normal_(self.cls_token, std=0.02)
        nn.init.normal_(self.position, std=0.02)

    def patch_tokens(self, images: Tensor) -> Tensor:
        return self.patch_embed(images)

    def prepend_tokens(self, patches: Tensor) -> Tensor:
        assert self.extra_tokens == 1
        cls = self.cls_token.expand(patches.shape[0], -1, -1)
        tokens = torch.cat((cls, patches), dim=1)
        assert tokens.shape[1] == self.config.patch_tokens + 1
        return tokens

    def add_positions(self, tokens: Tensor) -> Tensor:
        assert tokens.shape[1] == self.position.shape[1]
        assert tokens.shape[2] == self.config.width
        return tokens + self.position

    def encode_tokens(self, tokens: Tensor) -> Tensor:
        for block in self.blocks:
            tokens = block(tokens)
            assert tokens.shape[2] == self.config.width
        return self.norm(tokens)

    def forward_features(self, images: Tensor) -> tuple[Tensor, Tensor]:
        patches = self.patch_tokens(images)
        tokens = self.prepend_tokens(patches)
        tokens = self.add_positions(tokens)
        tokens = self.encode_tokens(tokens)
        cls = tokens[:, 0]
        assert cls.shape == (images.shape[0], self.config.width)
        return tokens, cls

    def forward(self, images: Tensor) -> Tensor:
        _, cls = self.forward_features(images)
        logits = self.head(cls)
        assert logits.shape == (images.shape[0], self.config.classes)
        return logits


class DistilledViT(TinyViT):
    def __init__(self, config: ToyConfig):
        super().__init__(config, extra_tokens=2)
        self.dist_token = nn.Parameter(torch.zeros(1, 1, config.width))
        self.dist_head = nn.Linear(config.width, config.classes)

    def prepend_tokens(self, patches: torch.Tensor) -> torch.Tensor:
        batch, patch_count, width = patches.shape
        assert patch_count == self.config.patch_tokens
        assert width == self.config.width
        cls = self.cls_token.expand(batch, -1, -1)
        dist = self.dist_token.expand(batch, -1, -1)
        tokens = torch.cat((cls, dist, patches), dim=1)
        assert tokens.shape == (
            batch,
            self.config.patch_tokens + 2,
            self.config.width,
        )
        return tokens

    def forward(
        self, images: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor] | torch.Tensor:
        encoded, cls = super().forward_features(images)
        class_logits = self.head(cls)
        distill_logits = self.dist_head(encoded[:, 1])
        if self.training:
            return class_logits, distill_logits
        return (class_logits + distill_logits) / 2


ShapeReport = dict[str, tuple[int, ...]]
Demo = Callable[[ToyConfig], ShapeReport]


def synthetic_images(config: ToyConfig) -> Tensor:
    values = torch.linspace(
        -1.0,
        1.0,
        steps=(
            config.batch_size
            * config.channels
            * config.image_size
            * config.image_size
        ),
    )
    return values.reshape(
        config.batch_size,
        config.channels,
        config.image_size,
        config.image_size,
    )


def run_vit_demo(config: ToyConfig) -> ShapeReport:
    model = TinyViT(config).cpu()
    images = synthetic_images(config)
    patches = model.patch_tokens(images)
    tokens, cls = model.forward_features(images)
    logits = model.head(cls)
    labels = torch.arange(config.batch_size) % config.classes
    loss = F.cross_entropy(logits, labels)
    loss.backward()

    assert images.device.type == "cpu"
    assert patches.shape == (
        config.batch_size,
        config.patch_tokens,
        config.width,
    )
    assert tokens.shape == (
        config.batch_size,
        config.patch_tokens + 1,
        config.width,
    )
    assert cls.shape == (config.batch_size, config.width)
    assert logits.shape == (config.batch_size, config.classes)
    assert model.patch_embed.proj.weight.grad is not None
    return {
        "images": tuple(images.shape),
        "patches": tuple(patches.shape),
        "tokens": tuple(tokens.shape),
        "cls": tuple(cls.shape),
        "logits": tuple(logits.shape),
    }


def deit_loss(
    class_logits: torch.Tensor,
    distill_logits: torch.Tensor,
    labels: torch.Tensor,
    teacher_logits: torch.Tensor,
    hard: bool,
    alpha: float,
    temperature: float,
) -> torch.Tensor:
    if (
        isinstance(alpha, bool)
        or not isinstance(alpha, (int, float))
        or not math.isfinite(alpha)
        or not 0 <= alpha <= 1
    ):
        raise ValueError("alpha must be a finite number between 0 and 1")
    if (
        isinstance(temperature, bool)
        or not isinstance(temperature, (int, float))
        or not math.isfinite(temperature)
        or temperature <= 0
    ):
        raise ValueError("temperature must be a positive finite number")

    teacher_logits = teacher_logits.detach()
    class_loss = F.cross_entropy(class_logits, labels)
    if hard:
        teacher_targets = teacher_logits.argmax(dim=-1)
        distill_loss = F.cross_entropy(distill_logits, teacher_targets)
    else:
        distill_loss = F.kl_div(
            F.log_softmax(distill_logits / temperature, dim=-1),
            F.log_softmax(teacher_logits / temperature, dim=-1),
            reduction="sum",
            log_target=True,
        )
        distill_loss = (
            distill_loss * temperature**2 / distill_logits.numel()
        )
    return (1 - alpha) * class_loss + alpha * distill_loss


def run_deit_demo(config: ToyConfig) -> ShapeReport:
    model = DistilledViT(config)
    images = torch.randn(
        config.batch_size,
        config.channels,
        config.image_size,
        config.image_size,
    )
    labels = torch.arange(config.batch_size) % config.classes
    teacher_logits = torch.linspace(
        -2.0,
        2.0,
        steps=config.batch_size * config.classes,
    ).reshape(config.batch_size, config.classes)
    teacher_logits.requires_grad_()
    model.train()
    training_output = model(images)
    assert isinstance(training_output, tuple)
    class_logits, distill_logits = training_output
    hard_loss = deit_loss(
        class_logits,
        distill_logits,
        labels,
        teacher_logits,
        hard=True,
        alpha=0.5,
        temperature=1.0,
    )
    soft_loss = deit_loss(
        class_logits,
        distill_logits,
        labels,
        teacher_logits,
        hard=False,
        alpha=0.5,
        temperature=2.0,
    )
    (hard_loss + soft_loss).backward()
    assert model.dist_token.grad is not None
    assert teacher_logits.grad is None

    model.eval()
    with torch.no_grad():
        fused_logits = model(images)
        encoded, cls = model.forward_features(images)
        expected_fused = (
            model.head(cls) + model.dist_head(encoded[:, 1])
        ) / 2
    assert torch.allclose(fused_logits, expected_fused)
    return {
        "images": tuple(images.shape),
        "tokens": (
            config.batch_size,
            config.patch_tokens + 2,
            config.width,
        ),
        "class_logits": tuple(class_logits.shape),
        "distill_logits": tuple(distill_logits.shape),
    }


BLIP2_TOY_QUERY_TOKENS = 8


class LearnedQueryResampler(nn.Module):
    def __init__(self, width: int, heads: int, query_tokens: int) -> None:
        super().__init__()
        for name, value in (
            ("width", width),
            ("heads", heads),
            ("query_tokens", query_tokens),
        ):
            if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
                raise ValueError(f"{name} must be a positive integer")
        if width % heads != 0:
            raise ValueError("width must be divisible by heads")
        self.width = width
        self.queries = nn.Parameter(torch.randn(1, query_tokens, width) * 0.02)
        self.cross_attention = nn.MultiheadAttention(width, heads, batch_first=True)

    def forward(self, image_tokens: Tensor) -> tuple[Tensor, Tensor]:
        if not isinstance(image_tokens, Tensor):
            raise TypeError("image_tokens must be a tensor")
        if image_tokens.ndim != 3:
            raise ValueError("image_tokens must be a rank-3 tensor")
        if not torch.is_floating_point(image_tokens):
            raise TypeError("image_tokens must use a floating dtype")
        if image_tokens.shape[0] == 0 or image_tokens.shape[1] == 0:
            raise ValueError("image_tokens must have non-empty batch and sequence")
        if image_tokens.shape[2] != self.width:
            raise ValueError("image_tokens width must match the resampler")
        if image_tokens.dtype != self.queries.dtype:
            raise TypeError("image_tokens dtype must match the resampler")
        if image_tokens.device != self.queries.device:
            raise ValueError("image_tokens device must match the resampler")
        if not bool(torch.isfinite(image_tokens).all().item()):
            raise ValueError("image_tokens must be finite")
        queries = self.queries.expand(image_tokens.shape[0], -1, -1)
        output, weights = self.cross_attention(
            queries,
            image_tokens,
            image_tokens,
            need_weights=True,
            average_attn_weights=True,
        )
        assert output.shape == (
            image_tokens.shape[0],
            self.queries.shape[1],
            self.width,
        )
        assert weights.shape == (
            image_tokens.shape[0],
            self.queries.shape[1],
            image_tokens.shape[1],
        )
        return output, weights


DEMOS = {'vit': run_vit_demo}


def window_partition(features: torch.Tensor, window_size: int) -> torch.Tensor:
    batch, height, width, channels = features.shape
    assert height % window_size == 0
    assert width % window_size == 0
    windows = (
        features.view(
            batch,
            height // window_size,
            window_size,
            width // window_size,
            window_size,
            channels,
        )
        .permute(0, 1, 3, 2, 4, 5)
        .reshape(-1, window_size * window_size, channels)
    )
    return windows


def window_reverse(
    windows: torch.Tensor,
    window_size: int,
    height: int,
    width: int,
    batch: int,
) -> torch.Tensor:
    channels = windows.shape[-1]
    features = (
        windows.view(
            batch,
            height // window_size,
            width // window_size,
            window_size,
            window_size,
            channels,
        )
        .permute(0, 1, 3, 2, 4, 5)
        .reshape(batch, height, width, channels)
    )
    return features


def shifted_window_mask(
    height: int,
    width: int,
    window_size: int,
    shift_size: int,
    device: torch.device,
) -> torch.Tensor:
    labels = torch.zeros((1, height, width, 1), device=device)
    height_slices = (
        slice(0, -window_size),
        slice(-window_size, -shift_size),
        slice(-shift_size, None),
    )
    width_slices = height_slices
    region = 0
    for height_slice in height_slices:
        for width_slice in width_slices:
            labels[:, height_slice, width_slice, :] = region
            region += 1
    label_windows = window_partition(labels, window_size).squeeze(-1)
    differences = label_windows.unsqueeze(1) - label_windows.unsqueeze(2)
    return differences.masked_fill(differences != 0, -100.0).masked_fill(
        differences == 0, 0.0
    )


class WindowAttention(MultiHeadSelfAttention):
    def forward(
        self,
        tokens: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        batch_windows, token_count, width = tokens.shape
        qkv = (
            self.qkv(tokens)
            .reshape(batch_windows, token_count, 3, self.heads, self.head_dim)
            .permute(2, 0, 3, 1, 4)
        )
        query, key, value = qkv.unbind(0)
        scores = (query @ key.transpose(-2, -1)) * self.scale
        if attention_mask is not None:
            window_count = attention_mask.shape[0]
            scores = scores.view(
                batch_windows // window_count,
                window_count,
                self.heads,
                token_count,
                token_count,
            )
            scores = scores + attention_mask[None, :, None, :, :]
            scores = scores.view(
                batch_windows, self.heads, token_count, token_count
            )
        mixed = scores.softmax(dim=-1) @ value
        mixed = mixed.transpose(1, 2).reshape(batch_windows, token_count, width)
        return self.proj(mixed)


def shifted_window_block(
    features: torch.Tensor,
    attention: WindowAttention,
    window_size: int,
    shift_size: int,
) -> torch.Tensor:
    batch, height, width, _ = features.shape
    shifted = torch.roll(
        features,
        shifts=(-shift_size, -shift_size),
        dims=(1, 2),
    )
    windows = window_partition(shifted, window_size)
    mask = shifted_window_mask(
        height,
        width,
        window_size,
        shift_size,
        features.device,
    )
    attended = attention(windows, mask)
    restored = window_reverse(
        attended,
        window_size,
        height,
        width,
        batch,
    )
    return torch.roll(restored, shifts=(shift_size, shift_size), dims=(1, 2))


def patch_merge(
    features: torch.Tensor,
    projection: nn.Linear,
) -> torch.Tensor:
    batch, height, width, channels = features.shape
    assert height % 2 == 0 and width % 2 == 0
    merged = torch.cat(
        (
            features[:, 0::2, 0::2],
            features[:, 1::2, 0::2],
            features[:, 0::2, 1::2],
            features[:, 1::2, 1::2],
        ),
        dim=-1,
    )
    assert merged.shape == (batch, height // 2, width // 2, 4 * channels)
    return projection(merged)


def run_swin_demo(config: ToyConfig) -> ShapeReport:
    assert config.grid_size % 4 == 0
    patch_embed = PatchEmbedding(config)
    attention = WindowAttention(config.width, config.heads)
    merge_projection = nn.Linear(4 * config.width, 2 * config.width)
    images = torch.randn(
        config.batch_size,
        config.channels,
        config.image_size,
        config.image_size,
    )
    patches = patch_embed(images)
    features = patches.view(
        config.batch_size,
        config.grid_size,
        config.grid_size,
        config.width,
    )
    windows = window_partition(features, 4)
    partitioned_then_reversed = window_reverse(
        windows,
        4,
        config.grid_size,
        config.grid_size,
        config.batch_size,
    )
    assert torch.equal(partitioned_then_reversed, features)
    mask = shifted_window_mask(
        config.grid_size,
        config.grid_size,
        4,
        2,
        features.device,
    )
    assert torch.any(mask == 0)
    assert torch.any(mask < 0)
    ordinary_windows = attention(windows)
    ordinary = window_reverse(
        ordinary_windows,
        4,
        config.grid_size,
        config.grid_size,
        config.batch_size,
    )
    shifted = shifted_window_block(ordinary, attention, 4, 2)
    merged = patch_merge(shifted, merge_projection)
    merged.square().mean().backward()
    assert merge_projection.weight.grad is not None
    return {
        "features": tuple(features.shape),
        "windows": tuple(windows.shape),
        "shifted": tuple(shifted.shape),
        "merged": tuple(merged.shape),
    }


DEMOS["deit"] = run_deit_demo


DEMOS["swin"] = run_swin_demo


class SpatialReductionAttention(nn.Module):
    def __init__(
        self,
        width: int,
        heads: int,
        reduction_ratio: int,
    ):
        super().__init__()
        assert width % heads == 0
        assert reduction_ratio >= 1
        self.width = width
        self.heads = heads
        self.head_dim = width // heads
        self.scale = self.head_dim**-0.5
        self.reduction_ratio = reduction_ratio
        self.query = nn.Linear(width, width)
        self.reduction = nn.Conv2d(
            width,
            width,
            kernel_size=reduction_ratio,
            stride=reduction_ratio,
        )
        self.norm = nn.LayerNorm(width)
        self.key_value = nn.Linear(width, 2 * width)
        self.output = nn.Linear(width, width)

    def forward(
        self, features: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor]:
        batch, height, width_tokens, channels = features.shape
        query_tokens = features.reshape(batch, height * width_tokens, channels)
        query = (
            self.query(query_tokens)
            .reshape(batch, height * width_tokens, self.heads, self.head_dim)
            .transpose(1, 2)
        )
        reduced_map = self.reduction(features.permute(0, 3, 1, 2))
        reduced_tokens = reduced_map.flatten(2).transpose(1, 2)
        reduced_tokens = self.norm(reduced_tokens)
        key_value = (
            self.key_value(reduced_tokens)
            .reshape(batch, reduced_tokens.shape[1], 2, self.heads, self.head_dim)
            .permute(2, 0, 3, 1, 4)
        )
        key, value = key_value.unbind(0)
        attention = (query @ key.transpose(-2, -1) * self.scale).softmax(dim=-1)
        mixed = attention @ value
        mixed = mixed.transpose(1, 2).reshape(
            batch, height * width_tokens, channels
        )
        return self.output(mixed), reduced_tokens


def run_pvt_demo(config: ToyConfig) -> ShapeReport:
    patch_embed = PatchEmbedding(config)
    attention = SpatialReductionAttention(config.width, config.heads, 4)
    images = torch.randn(
        config.batch_size,
        config.channels,
        config.image_size,
        config.image_size,
    )
    patches = patch_embed(images)
    features = patches.view(
        config.batch_size,
        config.grid_size,
        config.grid_size,
        config.width,
    )
    output, reduced = attention(features)
    output.square().mean().backward()
    assert attention.query.weight.grad is not None
    return {
        "features": tuple(features.shape),
        "queries": tuple(output.shape),
        "key_values": tuple(reduced.shape),
    }


DEMOS["pvt"] = run_pvt_demo


class PoolingAttention(nn.Module):
    def __init__(
        self,
        width: int,
        heads: int,
        query_stride: int,
        key_value_stride: int,
    ):
        super().__init__()
        assert width % heads == 0
        assert query_stride >= 1
        assert key_value_stride >= 1
        self.width = width
        self.heads = heads
        self.head_dim = width // heads
        self.scale = self.head_dim**-0.5
        self.query_stride = query_stride
        self.key_value_stride = key_value_stride
        self.query = nn.Linear(width, width)
        self.key = nn.Linear(width, width)
        self.value = nn.Linear(width, width)
        self.output = nn.Linear(width, width)

    @staticmethod
    def _pool(features: Tensor, stride: int) -> Tensor:
        pooled = F.avg_pool2d(
            features.permute(0, 3, 1, 2),
            kernel_size=stride,
            stride=stride,
        )
        return pooled.permute(0, 2, 3, 1)

    def _heads(self, tokens: Tensor) -> Tensor:
        batch, token_count, _ = tokens.shape
        return tokens.reshape(
            batch, token_count, self.heads, self.head_dim
        ).transpose(1, 2)

    def forward(self, features: Tensor) -> tuple[Tensor, Tensor, Tensor]:
        batch, height, width_tokens, channels = features.shape
        input_tokens = features.reshape(batch, height * width_tokens, channels)
        query_map = self.query(input_tokens).view(
            batch, height, width_tokens, channels
        )
        key_map = self.key(input_tokens).view(batch, height, width_tokens, channels)
        value_map = self.value(input_tokens).view(
            batch, height, width_tokens, channels
        )
        query_map = self._pool(query_map, self.query_stride)
        key_map = self._pool(key_map, self.key_value_stride)
        value_map = self._pool(value_map, self.key_value_stride)
        query_tokens = query_map.reshape(batch, -1, channels)
        key_tokens = key_map.reshape(batch, -1, channels)
        value_tokens = value_map.reshape(batch, -1, channels)
        query = self._heads(query_tokens)
        key = self._heads(key_tokens)
        value = self._heads(value_tokens)
        attention = (query @ key.transpose(-2, -1) * self.scale).softmax(dim=-1)
        mixed = attention @ value
        mixed = mixed.transpose(1, 2).reshape(batch, query_tokens.shape[1], channels)
        output_tokens = self.output(mixed)
        output_height = height // self.query_stride
        output_width = width_tokens // self.query_stride
        output = output_tokens.view(batch, output_height, output_width, channels)
        return output, query_tokens, key_tokens


def run_mvit_demo(config: ToyConfig) -> ShapeReport:
    patch_embed = PatchEmbedding(config)
    attention = PoolingAttention(config.width, config.heads, 2, 4)
    images = torch.randn(
        config.batch_size,
        config.channels,
        config.image_size,
        config.image_size,
    )
    patches = patch_embed(images)
    features = patches.view(
        config.batch_size,
        config.grid_size,
        config.grid_size,
        config.width,
    )
    projection_input_counts: dict[str, int] = {}

    def record_projection_input(name: str):
        def hook(_module: nn.Module, inputs: tuple[Tensor, ...]) -> None:
            projection_input_counts[name] = inputs[0].shape[1]

        return hook

    hooks = [
        projection.register_forward_pre_hook(record_projection_input(name))
        for name, projection in (
            ("query", attention.query),
            ("key", attention.key),
            ("value", attention.value),
        )
    ]
    output, queries, key_values = attention(features)
    for hook in hooks:
        hook.remove()
    assert projection_input_counts == {
        "query": config.patch_tokens,
        "key": config.patch_tokens,
        "value": config.patch_tokens,
    }
    output.square().mean().backward()
    assert attention.query.weight.grad is not None
    return {
        "features": tuple(features.shape),
        "queries": tuple(queries.shape),
        "key_values": tuple(key_values.shape),
        "output": tuple(output.shape),
    }


DEMOS["mvit"] = run_mvit_demo


def _positive_integer(name: str, value: int) -> None:
    if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
        raise ValueError(f"{name} must be a positive integer")


def _mask_ratio(ratio: float) -> None:
    if (
        isinstance(ratio, bool)
        or not isinstance(ratio, (int, float))
        or not math.isfinite(ratio)
        or not 0 < ratio < 1
    ):
        raise ValueError("ratio must be a finite number with 0 < ratio < 1")


def mask_partition_sizes(
    patch_count: int,
    ratio: float,
) -> tuple[int, int]:
    _positive_integer("patch_count", patch_count)
    _mask_ratio(ratio)
    if patch_count < 2:
        raise ValueError("patch_count must be at least 2")
    masked_count = math.floor(patch_count * ratio + 0.5)
    masked_count = max(1, min(patch_count - 1, masked_count))
    return patch_count - masked_count, masked_count


def patchify(images: torch.Tensor, patch_size: int) -> torch.Tensor:
    if not isinstance(images, torch.Tensor) or images.ndim != 4:
        raise ValueError("images must be a rank-4 torch.Tensor")
    _positive_integer("patch_size", patch_size)
    batch, channels, height, width = images.shape
    if any(axis <= 0 for axis in (batch, channels, height, width)):
        raise ValueError("images must have non-zero batch, channel, and spatial axes")
    if height % patch_size != 0 or width % patch_size != 0:
        raise ValueError("image height and width must be divisible by patch_size")
    rows, columns = height // patch_size, width // patch_size
    return (
        images.reshape(
            batch,
            channels,
            rows,
            patch_size,
            columns,
            patch_size,
        )
        .permute(0, 2, 4, 3, 5, 1)
        .reshape(batch, rows * columns, -1)
    )


def random_mask(
    patches: torch.Tensor,
    ratio: float,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    if not isinstance(patches, torch.Tensor) or patches.ndim != 3:
        raise ValueError("patches must be a rank-3 torch.Tensor")
    batch, count, width = patches.shape
    if batch == 0 or width == 0:
        raise ValueError("patches must contain non-empty patch vectors")
    keep, _ = mask_partition_sizes(count, ratio)
    shuffle = torch.rand(batch, count, device=patches.device).argsort(1)
    restore, ids = shuffle.argsort(1), shuffle[:, :keep]
    visible = torch.gather(
        patches,
        1,
        ids.unsqueeze(-1).expand(-1, -1, width),
    )
    mask = torch.ones(batch, count, dtype=torch.bool, device=patches.device)
    mask[:, :keep] = False
    return visible, ids, torch.gather(mask, 1, restore)


def _validate_patch_mask(
    mask: torch.Tensor,
    batch: int,
    patch_count: int,
    device: torch.device,
) -> None:
    if batch <= 0 or patch_count < 2:
        raise ValueError(
            "mask batches must be non-empty and contain at least 2 patches"
        )
    expected_shape = (batch, patch_count)
    if not isinstance(mask, torch.Tensor) or tuple(mask.shape) != expected_shape:
        raise ValueError(
            f"mask shape must be {expected_shape}, got "
            f"{getattr(mask, 'shape', None)}"
        )
    if mask.dtype != torch.bool:
        raise ValueError("mask must be a boolean tensor")
    if mask.device != device:
        raise ValueError("mask must be on the same device as patch tokens")
    selected = mask.sum(dim=1)
    invalid = (selected == 0) | (selected == patch_count)
    if bool(invalid.any().item()):
        raise ValueError(
            "mask must select at least one masked and one visible patch per sample"
        )


def restore_mae_decoder_tokens(
    decoder_cls: torch.Tensor,
    visible_tokens: torch.Tensor,
    visible_ids: torch.Tensor,
    mask_token: torch.Tensor,
    decoder_position: torch.Tensor,
) -> torch.Tensor:
    if not isinstance(decoder_cls, torch.Tensor) or decoder_cls.ndim != 3:
        raise ValueError("decoder_cls must be a rank-3 torch.Tensor")
    batch, cls_count, width = decoder_cls.shape
    if batch == 0 or cls_count != 1 or width == 0:
        raise ValueError("decoder_cls must have shape (batch, 1, decoder width)")
    if not isinstance(visible_tokens, torch.Tensor) or visible_tokens.ndim != 3:
        raise ValueError("visible_tokens must be a rank-3 torch.Tensor")
    visible_batch, visible_count, visible_width = visible_tokens.shape
    if visible_batch != batch or visible_width != width:
        raise ValueError(
            "visible_tokens must match decoder_cls batch and decoder width"
        )
    if not isinstance(decoder_position, torch.Tensor) or decoder_position.ndim != 3:
        raise ValueError("decoder_position must be a rank-3 torch.Tensor")
    if decoder_position.shape[0] != 1 or decoder_position.shape[2] != width:
        raise ValueError(
            "decoder_position must have shape (1, patch count + 1, decoder width)"
        )
    patch_count = decoder_position.shape[1] - 1
    if patch_count < 2:
        raise ValueError("decoder_position must describe at least 2 patches")
    if not 1 <= visible_count < patch_count:
        raise ValueError(
            "visible_tokens must leave at least one visible and one masked patch"
        )
    if not isinstance(visible_ids, torch.Tensor) or visible_ids.shape != (
        batch,
        visible_count,
    ):
        raise ValueError("visible_ids must match visible_tokens batch and count")
    if visible_ids.dtype != torch.long:
        raise ValueError("visible_ids must be a torch.long tensor")
    if not isinstance(mask_token, torch.Tensor) or mask_token.shape != (1, 1, width):
        raise ValueError("mask_token must have shape (1, 1, decoder width)")

    floating_tensors = (
        decoder_cls,
        visible_tokens,
        mask_token,
        decoder_position,
    )
    if not all(torch.is_floating_point(tensor) for tensor in floating_tensors):
        raise ValueError("decoder tokens and positions must be floating-point tensors")
    if any(tensor.dtype != decoder_cls.dtype for tensor in floating_tensors[1:]):
        raise ValueError("decoder tokens and positions must use the same dtype")
    if any(tensor.device != decoder_cls.device for tensor in floating_tensors[1:]):
        raise ValueError("decoder tokens and positions must use the same device")
    if visible_ids.device != decoder_cls.device:
        raise ValueError("visible_ids must use the decoder token device")
    if visible_ids.min().item() < 0 or visible_ids.max().item() >= patch_count:
        raise ValueError("visible_ids must refer to original patch slots")
    if visible_count > 1:
        ordered_ids = visible_ids.sort(dim=1).values
        if bool((ordered_ids[:, 1:] == ordered_ids[:, :-1]).any().item()):
            raise ValueError("visible_ids must be unique within each sample")

    restored_patches = torch.scatter(
        mask_token.expand(batch, patch_count, width),
        1,
        visible_ids.unsqueeze(-1).expand(-1, -1, width),
        visible_tokens,
    )
    restored = torch.cat((decoder_cls, restored_patches), dim=1)
    return restored + decoder_position


def mae_masked_loss(
    prediction: torch.Tensor,
    targets: torch.Tensor,
    mask: torch.Tensor,
) -> torch.Tensor:
    if not isinstance(prediction, torch.Tensor) or prediction.ndim != 3:
        raise ValueError("prediction must be a rank-3 torch.Tensor")
    if not isinstance(targets, torch.Tensor) or targets.shape != prediction.shape:
        raise ValueError("targets must match the prediction shape")
    if not torch.is_floating_point(prediction) or not torch.is_floating_point(
        targets
    ):
        raise ValueError("prediction and targets must be floating-point tensors")
    if prediction.device != targets.device or prediction.dtype != targets.dtype:
        raise ValueError("prediction and targets must share device and dtype")
    batch, patch_count, target_width = prediction.shape
    if batch == 0 or patch_count < 2 or target_width == 0:
        raise ValueError(
            "prediction must contain non-empty batches, patches, and patch targets"
        )
    _validate_patch_mask(mask, batch, patch_count, prediction.device)
    selected_prediction = prediction[mask]
    selected_targets = targets[mask]
    return (selected_prediction - selected_targets).square().mean()


def beit_masked_loss(
    logits: torch.Tensor,
    targets: torch.Tensor,
    mask: torch.Tensor,
) -> torch.Tensor:
    if not isinstance(logits, torch.Tensor) or logits.ndim != 3:
        raise ValueError("logits must be a rank-3 torch.Tensor")
    if not torch.is_floating_point(logits) or logits.shape[2] == 0:
        raise ValueError("logits must contain floating-point vocabulary scores")
    batch, patch_count, vocabulary_size = logits.shape
    if batch == 0 or patch_count < 2:
        raise ValueError("logits must contain non-empty batches and patches")
    if not isinstance(targets, torch.Tensor) or targets.shape != (
        batch,
        patch_count,
    ):
        raise ValueError("targets must match the logits batch and patch axes")
    if targets.dtype != torch.long:
        raise ValueError("targets must be a torch.long tensor")
    if targets.device != logits.device:
        raise ValueError("targets and logits must use the same device")
    _validate_patch_mask(mask, batch, patch_count, logits.device)
    if targets.min().item() < 0 or targets.max().item() >= vocabulary_size:
        raise ValueError("targets must fall inside the logits vocabulary")
    return F.cross_entropy(logits[mask], targets[mask])


def synthetic_visual_targets(
    images: torch.Tensor,
    patch_size: int,
    vocabulary_size: int,
) -> torch.Tensor:
    _positive_integer("vocabulary_size", vocabulary_size)
    return (
        (patchify(images, patch_size).mean(-1) * 1000)
        .round()
        .long()
        .abs()
        .remainder(vocabulary_size)
    )


def encode_with_patch_mask(
    backbone: TinyViT,
    images: torch.Tensor,
    mask: torch.Tensor,
    mask_token: torch.Tensor,
) -> torch.Tensor:
    patches = backbone.patch_tokens(images)
    _validate_patch_mask(
        mask,
        patches.shape[0],
        patches.shape[1],
        patches.device,
    )
    if not isinstance(mask_token, torch.Tensor) or mask_token.shape != (
        1,
        1,
        patches.shape[-1],
    ):
        raise ValueError("mask_token must have shape (1, 1, backbone width)")
    if mask_token.device != patches.device:
        raise ValueError("mask_token must be on the same device as patch tokens")
    replacement = mask_token.expand(images.shape[0], patches.shape[1], -1)
    tokens = backbone.prepend_tokens(
        torch.where(mask.unsqueeze(-1), replacement, patches)
    )
    return backbone.encode_tokens(backbone.add_positions(tokens))


class ToyMAE(nn.Module):
    def __init__(self, config: ToyConfig, mask_ratio: float = 0.75):
        super().__init__()
        visible_patches, masked_patches = mask_partition_sizes(
            config.patch_tokens,
            mask_ratio,
        )
        self.config = config
        self.mask_ratio = mask_ratio
        self.visible_patches = visible_patches
        self.masked_patches = masked_patches
        self.backbone = TinyViT(config)
        self.decoder_width = 32
        self.encoder_to_decoder = nn.Linear(config.width, self.decoder_width)
        self.mask_token = nn.Parameter(torch.zeros(1, 1, self.decoder_width))
        self.decoder_position = nn.Parameter(
            torch.zeros(1, config.patch_tokens + 1, self.decoder_width)
        )
        self.decoder = EncoderBlock(self.decoder_width, 4, 2)
        self.pixel_head = nn.Linear(
            self.decoder_width,
            config.patch_size * config.patch_size * config.channels,
        )
        nn.init.normal_(self.mask_token, std=0.02)
        nn.init.normal_(self.decoder_position, std=0.02)

    def forward(
        self, images: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        patches = self.backbone.patch_tokens(images)
        visible, ids, mask = random_mask(patches, self.mask_ratio)
        batch = images.shape[0]
        expected_visible = self.visible_patches
        assert visible.shape == (batch, expected_visible, self.config.width)
        assert torch.equal(
            mask.sum(dim=1),
            torch.full(
                (batch,),
                self.config.patch_tokens - expected_visible,
                device=mask.device,
            ),
        )
        patch_positions = self.backbone.position[:, 1:].expand(batch, -1, -1)
        visible_positions = torch.gather(
            patch_positions,
            1,
            ids.unsqueeze(-1).expand(-1, -1, self.config.width),
        )
        cls = (
            self.backbone.cls_token + self.backbone.position[:, :1]
        ).expand(batch, -1, -1)
        encoder_tokens = self.backbone.encode_tokens(
            torch.cat((cls, visible + visible_positions), dim=1)
        )
        visible_decoder = self.encoder_to_decoder(encoder_tokens[:, 1:])
        decoder_cls = self.encoder_to_decoder(encoder_tokens[:, :1])
        decoder_tokens = restore_mae_decoder_tokens(
            decoder_cls,
            visible_decoder,
            ids,
            self.mask_token,
            self.decoder_position,
        )
        decoded = self.decoder(decoder_tokens)
        prediction = self.pixel_head(decoded[:, 1:])
        assert encoder_tokens.shape == (
            batch,
            expected_visible + 1,
            self.config.width,
        )
        assert prediction.shape == (
            batch,
            self.config.patch_tokens,
            self.config.patch_size
            * self.config.patch_size
            * self.config.channels,
        )
        return encoder_tokens, prediction, mask


class ToyBEiT(nn.Module):
    def __init__(self, config: ToyConfig, vocabulary_size: int = 32):
        super().__init__()
        _positive_integer("vocabulary_size", vocabulary_size)
        mask_partition_sizes(config.patch_tokens, 0.5)
        self.config = config
        self.vocabulary_size = vocabulary_size
        self.backbone = TinyViT(config)
        self.mask_token = nn.Parameter(torch.zeros(1, 1, config.width))
        self.vocabulary_head = nn.Linear(config.width, vocabulary_size)

    def forward(
        self,
        images: torch.Tensor,
        mask: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        encoded = encode_with_patch_mask(
            self.backbone,
            images,
            mask,
            self.mask_token,
        )
        logits = self.vocabulary_head(encoded[:, 1:])
        targets = synthetic_visual_targets(
            images,
            self.config.patch_size,
            self.vocabulary_size,
        )
        assert encoded.shape == (
            images.shape[0],
            self.config.patch_tokens + 1,
            self.config.width,
        )
        assert logits.shape == (
            images.shape[0],
            self.config.patch_tokens,
            self.vocabulary_size,
        )
        assert targets.shape == (images.shape[0], self.config.patch_tokens)
        assert targets.min() >= 0
        assert targets.max() < self.vocabulary_size
        return encoded, logits, targets


def run_mae_demo(config: ToyConfig) -> ShapeReport:
    model = ToyMAE(config).cpu()
    images = synthetic_images(config)
    encoder_tokens, prediction, mask = model(images)
    targets = patchify(images, config.patch_size)
    expected_visible, expected_masked = mask_partition_sizes(
        config.patch_tokens,
        model.mask_ratio,
    )
    assert encoder_tokens.shape == (
        config.batch_size,
        expected_visible + 1,
        config.width,
    )
    assert prediction.shape == targets.shape
    assert mask.shape == (config.batch_size, config.patch_tokens)
    assert mask.dtype == torch.bool
    assert torch.equal(
        mask.sum(dim=1),
        torch.full((config.batch_size,), expected_masked),
    )
    assert expected_masked > 0
    decoder_patch_positions = model.decoder_position.detach()[0, 1:]
    assert (
        torch.unique(decoder_patch_positions, dim=0).shape[0]
        == config.patch_tokens
    )
    for sample_prediction, sample_mask in zip(prediction, mask):
        masked_predictions = sample_prediction[sample_mask]
        assert masked_predictions.shape[0] == expected_masked
        if expected_masked > 1:
            assert not torch.allclose(
                masked_predictions,
                masked_predictions[:1].expand_as(masked_predictions),
            )
    loss = mae_masked_loss(prediction, targets, mask)
    assert torch.isfinite(loss)
    loss.backward()
    gradient = model.backbone.patch_embed.proj.weight.grad
    assert gradient is not None
    assert torch.isfinite(gradient).all()
    assert torch.count_nonzero(gradient) > 0
    decoder_position_gradient = model.decoder_position.grad
    assert decoder_position_gradient is not None
    assert torch.isfinite(decoder_position_gradient).all()
    assert torch.count_nonzero(decoder_position_gradient) > 0
    return {
        "encoder_tokens": tuple(encoder_tokens.shape),
        "prediction": tuple(prediction.shape),
        "mask": tuple(mask.shape),
    }


def run_beit_demo(config: ToyConfig) -> ShapeReport:
    model = ToyBEiT(config).cpu()
    images = synthetic_images(config)
    mask_ratio = 0.4
    with torch.no_grad():
        patches = model.backbone.patch_tokens(images)
        _, _, mask = random_mask(patches, mask_ratio)
    expected_visible, expected_masked = mask_partition_sizes(
        config.patch_tokens,
        mask_ratio,
    )
    assert mask.shape == (config.batch_size, config.patch_tokens)
    assert mask.dtype == torch.bool
    assert torch.equal(
        mask.sum(dim=1),
        torch.full((config.batch_size,), expected_masked),
    )
    assert expected_masked > 0
    tokens, logits, targets = model(images, mask)
    assert tokens.shape == (
        config.batch_size,
        config.patch_tokens + 1,
        config.width,
    )
    assert logits.shape == (
        config.batch_size,
        config.patch_tokens,
        model.vocabulary_size,
    )
    assert targets.shape == (config.batch_size, config.patch_tokens)
    assert targets.min() >= 0
    assert targets.max() < model.vocabulary_size
    loss = beit_masked_loss(logits, targets, mask)
    assert torch.isfinite(loss)
    loss.backward()
    gradient = model.backbone.patch_embed.proj.weight.grad
    assert gradient is not None
    assert torch.isfinite(gradient).all()
    assert torch.count_nonzero(gradient) > 0
    return {
        "tokens": tuple(tokens.shape),
        "logits": tuple(logits.shape),
        "targets": tuple(targets.shape),
        "mask": tuple(mask.shape),
    }


DEMOS["mae"] = run_mae_demo
DEMOS["beit"] = run_beit_demo


def _positive_finite(name: str, value: float) -> float:
    if (
        isinstance(value, bool)
        or not isinstance(value, (int, float))
        or not math.isfinite(value)
        or value <= 0
    ):
        raise ValueError(f"{name} must be a positive finite number")
    return float(value)


def _unit_interval(name: str, value: float) -> float:
    if (
        isinstance(value, bool)
        or not isinstance(value, (int, float))
        or not math.isfinite(value)
        or not 0 <= value <= 1
    ):
        raise ValueError(f"{name} must be a finite number between 0 and 1")
    return float(value)


def _validate_logit_views(
    name: str,
    values: list[torch.Tensor],
    rank: int,
    require_finite: bool = True,
) -> tuple[int, ...]:
    if not isinstance(values, (list, tuple)) or not values:
        raise ValueError(f"{name} views must be a non-empty list")
    shape: tuple[int, ...] | None = None
    dtype: torch.dtype | None = None
    device: torch.device | None = None
    for value in values:
        if not isinstance(value, torch.Tensor) or value.ndim != rank:
            raise ValueError(f"{name} views must contain rank-{rank} tensors")
        if not torch.is_floating_point(value):
            raise ValueError(f"{name} views must be floating-point tensors")
        if any(axis == 0 for axis in value.shape):
            raise ValueError(f"{name} views must have non-empty dimensions")
        if require_finite and not bool(torch.isfinite(value).all().item()):
            raise ValueError(f"{name} views must contain finite values")
        if shape is None:
            shape = tuple(value.shape)
            dtype = value.dtype
            device = value.device
        elif (
            tuple(value.shape) != shape
            or value.dtype != dtype
            or value.device != device
        ):
            raise ValueError(f"{name} views must share shape, dtype, and device")
    assert shape is not None
    return shape


def _validate_center(
    center: torch.Tensor,
    logit_shape: tuple[int, ...],
    dtype: torch.dtype,
    device: torch.device,
) -> None:
    expected_shape = (1,) * (len(logit_shape) - 1) + (logit_shape[-1],)
    if not isinstance(center, torch.Tensor) or tuple(center.shape) != expected_shape:
        raise ValueError(f"center shape must be {expected_shape}")
    if not torch.is_floating_point(center):
        raise ValueError("center must be a floating-point tensor")
    if center.dtype != dtype or center.device != device:
        raise ValueError("center must share teacher dtype and device")
    if not bool(torch.isfinite(center).all().item()):
        raise ValueError("center must contain finite values")


@torch.no_grad()
def centered_teacher_targets(
    teacher: list[torch.Tensor],
    center: torch.Tensor,
    teacher_temperature: float = 0.04,
    center_momentum: float = 0.9,
) -> list[torch.Tensor]:
    teacher_temperature = _positive_finite(
        "teacher_temperature", teacher_temperature
    )
    center_momentum = _unit_interval("center_momentum", center_momentum)
    if not isinstance(teacher, (list, tuple)) or not teacher:
        raise ValueError("teacher views must be a non-empty list")
    first = teacher[0]
    if not isinstance(first, torch.Tensor):
        raise ValueError("teacher views must contain tensors")
    shape = _validate_logit_views("teacher", teacher, first.ndim)
    if first.ndim not in (2, 3):
        raise ValueError("teacher views must contain rank-2 or rank-3 tensors")
    _validate_center(center, shape, first.dtype, first.device)

    old_center = center.detach().clone()
    targets = [
        ((value.detach() - old_center) / teacher_temperature).softmax(dim=-1)
        for value in teacher
    ]
    for target in targets:
        if not bool(torch.isfinite(target).all().item()):
            raise RuntimeError("teacher target probabilities must be finite")
        if not torch.allclose(
            target.sum(dim=-1),
            torch.ones_like(target[..., 0]),
        ):
            raise RuntimeError("teacher target probabilities must sum to one")

    combined = torch.cat([value.detach() for value in teacher], dim=0)
    reduction_dims = tuple(range(combined.ndim - 1))
    batch_mean = combined.mean(dim=reduction_dims, keepdim=True)
    center.mul_(center_momentum).add_(
        batch_mean,
        alpha=1 - center_momentum,
    )
    return targets


def _dino_pair_indices(
    student_count: int,
    teacher_count: int,
) -> list[tuple[int, int]]:
    return [
        (teacher_index, student_index)
        for teacher_index in range(teacher_count)
        for student_index in range(student_count)
        if student_index != teacher_index
    ]


def dino_cross_view_loss(
    student: list[torch.Tensor],
    teacher: list[torch.Tensor],
    center: torch.Tensor | None = None,
    student_temperature: float = 0.1,
    teacher_temperature: float = 0.04,
    center_momentum: float = 0.9,
) -> torch.Tensor:
    student_temperature = _positive_finite(
        "student_temperature", student_temperature
    )
    teacher_temperature = _positive_finite(
        "teacher_temperature", teacher_temperature
    )
    center_momentum = _unit_interval("center_momentum", center_momentum)
    student_shape = _validate_logit_views("student", student, 2)
    teacher_shape = _validate_logit_views("teacher", teacher, 2)
    if len(student) < len(teacher):
        raise ValueError("DINO needs at least as many student views as teacher views")
    if student_shape != teacher_shape:
        raise ValueError("student and teacher views must share their logits shape")
    if student[0].dtype != teacher[0].dtype or student[0].device != teacher[0].device:
        raise ValueError("student and teacher views must share dtype and device")
    if center is None:
        center = torch.zeros(
            1,
            teacher_shape[-1],
            dtype=teacher[0].dtype,
            device=teacher[0].device,
        )

    pairs = _dino_pair_indices(len(student), len(teacher))
    if not pairs:
        raise ValueError("DINO needs at least one non-aligned cross-view pair")
    targets = centered_teacher_targets(
        teacher,
        center,
        teacher_temperature=teacher_temperature,
        center_momentum=center_momentum,
    )
    losses = [
        -(
            targets[teacher_index]
            * (student[student_index] / student_temperature).log_softmax(dim=-1)
        )
        .sum(dim=-1)
        .mean()
        for teacher_index, student_index in pairs
    ]
    loss = torch.stack(losses).mean()
    if not bool(torch.isfinite(loss).item()):
        raise RuntimeError("DINO loss must be finite")
    return loss


@torch.no_grad()
def ema_update(teacher: nn.Module, student: nn.Module, momentum: float) -> None:
    momentum = _unit_interval("momentum", momentum)
    teacher_parameters = list(teacher.parameters())
    student_parameters = list(student.parameters())
    if len(teacher_parameters) != len(student_parameters):
        raise ValueError("teacher and student must have matching parameters")
    parameter_pairs = list(zip(teacher_parameters, student_parameters))
    for target, source in parameter_pairs:
        if target.shape != source.shape:
            raise ValueError("teacher and student parameter shapes must match")
    for target, source in parameter_pairs:
        target.mul_(momentum).add_(source, alpha=1 - momentum)


class DINOStudentTeacherToy(nn.Module):
    def __init__(
        self,
        config: ToyConfig,
        projection_size: int = 32,
        student_temperature: float = 0.1,
        teacher_temperature: float = 0.04,
        center_momentum: float = 0.9,
    ):
        super().__init__()
        _positive_integer("projection_size", projection_size)
        self.student_temperature = _positive_finite(
            "student_temperature", student_temperature
        )
        self.teacher_temperature = _positive_finite(
            "teacher_temperature", teacher_temperature
        )
        self.center_momentum = _unit_interval(
            "center_momentum", center_momentum
        )
        self.student = TinyViT(config)
        self.student.head = nn.Identity()
        self.teacher = deepcopy(self.student)
        self.student_head = nn.Linear(config.width, projection_size)
        self.teacher_head = deepcopy(self.student_head)
        self.register_buffer("center", torch.zeros(1, projection_size))
        for module in (self.teacher, self.teacher_head):
            for parameter in module.parameters():
                parameter.requires_grad_(False)

    def student_logits(self, images: torch.Tensor) -> torch.Tensor:
        _, cls = self.student.forward_features(images)
        return self.student_head(cls)

    @torch.no_grad()
    def teacher_logits(self, images: torch.Tensor) -> torch.Tensor:
        _, cls = self.teacher.forward_features(images)
        return self.teacher_head(cls)


def ibot_masked_loss(
    student: list[torch.Tensor],
    teacher: list[torch.Tensor],
    masks: list[torch.Tensor],
    center: torch.Tensor,
    student_temperature: float = 0.1,
    teacher_temperature: float = 0.04,
    center_momentum: float = 0.9,
) -> torch.Tensor:
    student_temperature = _positive_finite(
        "student_temperature", student_temperature
    )
    teacher_temperature = _positive_finite(
        "teacher_temperature", teacher_temperature
    )
    center_momentum = _unit_interval("center_momentum", center_momentum)
    student_shape = _validate_logit_views(
        "student patch",
        student,
        3,
        require_finite=False,
    )
    teacher_shape = _validate_logit_views(
        "teacher patch",
        teacher,
        3,
        require_finite=False,
    )
    if student_shape != teacher_shape:
        raise ValueError("student and teacher patch views must share shape")
    if len(student) != len(teacher):
        raise ValueError("student and teacher patch view counts must match")
    if not isinstance(masks, (list, tuple)) or not masks:
        raise ValueError("masks must be a non-empty list")
    if len(masks) != len(student):
        raise ValueError("masks must match the patch view count")
    if student[0].dtype != teacher[0].dtype or student[0].device != teacher[0].device:
        raise ValueError("student and teacher patch views must share dtype and device")
    for student_view, teacher_view, mask in zip(student, teacher, masks):
        _validate_patch_mask(
            mask,
            student_shape[0],
            student_shape[1],
            student[0].device,
        )
        if not bool(torch.isfinite(student_view[mask]).all().item()):
            raise ValueError("masked student patch logits must be finite")
        if not bool(torch.isfinite(teacher_view[mask]).all().item()):
            raise ValueError("masked teacher patch logits must be finite")

    _validate_center(center, teacher_shape, teacher[0].dtype, teacher[0].device)
    old_center = center.detach().clone()
    masked_teacher_views = [
        teacher_view.detach()[mask]
        for teacher_view, mask in zip(teacher, masks)
    ]
    masked_teacher = torch.cat(masked_teacher_views, dim=0)
    targets = []
    for teacher_view, mask in zip(teacher, masks):
        image_targets = []
        for sample_index in range(teacher_shape[0]):
            selected_teacher = teacher_view.detach()[
                sample_index,
                mask[sample_index],
            ]
            target = (
                (selected_teacher - old_center.reshape(1, -1))
                / teacher_temperature
            ).softmax(dim=-1)
            if not bool(torch.isfinite(target).all().item()):
                raise RuntimeError(
                    "masked teacher target probabilities must be finite"
                )
            if not torch.allclose(
                target.sum(dim=-1),
                torch.ones_like(target[:, 0]),
            ):
                raise RuntimeError(
                    "masked teacher target probabilities must sum to one"
                )
            image_targets.append(target)
        targets.append(image_targets)
    batch_mean = masked_teacher.mean(dim=0).reshape(1, 1, -1)
    with torch.no_grad():
        center.mul_(center_momentum).add_(
            batch_mean,
            alpha=1 - center_momentum,
        )
    view_losses = []
    for student_view, target_view, mask in zip(student, targets, masks):
        image_losses = []
        for sample_index in range(student_shape[0]):
            selected = mask[sample_index]
            student_log_probabilities = (
                student_view[sample_index, selected] / student_temperature
            ).log_softmax(dim=-1)
            image_losses.append(
                -(
                    target_view[sample_index] * student_log_probabilities
                )
                .sum(dim=-1)
                .mean()
            )
        view_losses.append(torch.stack(image_losses).mean())
    loss = torch.stack(view_losses).mean()
    if not bool(torch.isfinite(loss).item()):
        raise RuntimeError("iBOT loss must be finite")
    return loss


def koleo_loss(features: torch.Tensor) -> torch.Tensor:
    if not isinstance(features, torch.Tensor) or features.ndim != 2:
        raise ValueError("features must be a rank-2 tensor")
    if not torch.is_floating_point(features):
        raise ValueError("features must be floating-point")
    if features.shape[0] < 2:
        raise ValueError("features batch size must be at least 2")
    if features.shape[1] == 0:
        raise ValueError("features must have non-zero feature width")
    if not bool(torch.isfinite(features).all().item()):
        raise ValueError("features must contain finite values")
    values = F.normalize(features, dim=-1)
    distances = torch.cdist(values, values)
    eye = torch.eye(values.shape[0], dtype=torch.bool, device=values.device)
    nearest = distances.masked_fill(eye, float("inf")).min(dim=1).values
    loss = -torch.log(nearest.clamp_min(1e-6)).mean()
    if not bool(torch.isfinite(loss).item()):
        raise RuntimeError("KoLeo loss must be finite")
    return loss


def dinov2_koleo_loss(global_features: list[torch.Tensor]) -> torch.Tensor:
    """Average the CPU toy's two per-crop KoLeo terms.

    The paper Appendix describes KoLeo on the first global crop. Released code
    evaluates both global chunks separately and sums both terms into the
    optimization loss; only its logged metric averages them by dividing by two.
    This CPU toy keeps the same two-view selection but averages its two per-crop
    terms.
    """
    if not isinstance(global_features, (list, tuple)) or not global_features:
        raise ValueError("global feature views must be a non-empty list")
    return torch.stack([koleo_loss(features) for features in global_features]).mean()


class DINOv2Toy(nn.Module):
    def __init__(
        self,
        config: ToyConfig,
        projection_size: int = 32,
        student_temperature: float = 0.1,
        teacher_temperature: float = 0.04,
        center_momentum: float = 0.9,
    ):
        super().__init__()
        _positive_integer("projection_size", projection_size)
        self.student_temperature = _positive_finite(
            "student_temperature", student_temperature
        )
        self.teacher_temperature = _positive_finite(
            "teacher_temperature", teacher_temperature
        )
        self.center_momentum = _unit_interval(
            "center_momentum", center_momentum
        )
        self.student = TinyViT(config)
        self.student.head = nn.Identity()
        self.teacher = deepcopy(self.student)
        self.mask_token = nn.Parameter(torch.zeros(1, 1, config.width))
        self.student_cls_head = nn.Linear(config.width, projection_size)
        self.teacher_cls_head = deepcopy(self.student_cls_head)
        self.student_patch_head = nn.Linear(config.width, projection_size)
        self.teacher_patch_head = deepcopy(self.student_patch_head)
        self.register_buffer("cls_center", torch.zeros(1, projection_size))
        self.register_buffer(
            "patch_center", torch.zeros(1, 1, projection_size)
        )
        for module in (
            self.teacher,
            self.teacher_cls_head,
            self.teacher_patch_head,
        ):
            for parameter in module.parameters():
                parameter.requires_grad_(False)

    def forward(
        self,
        images: torch.Tensor,
        mask: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        student_tokens = encode_with_patch_mask(
            self.student,
            images,
            mask,
            self.mask_token,
        )
        with torch.no_grad():
            teacher_tokens, _ = self.teacher.forward_features(images)
            teacher_patch_logits = self.teacher_patch_head(
                teacher_tokens[:, 1:]
            )
        return (
            student_tokens[:, 0],
            teacher_tokens[:, 0],
            self.student_patch_head(student_tokens[:, 1:]),
            teacher_patch_logits,
        )


def _ssl_student_views(
    images: torch.Tensor,
    config: ToyConfig,
) -> list[torch.Tensor]:
    global_views = [images, images.flip(-1)]
    crop_size = max(config.patch_size, config.image_size // 2)
    crop_start = (config.image_size - crop_size) // 2
    local = F.interpolate(
        images[
            :,
            :,
            crop_start : crop_start + crop_size,
            crop_start : crop_start + crop_size,
        ],
        size=(config.image_size, config.image_size),
        mode="bilinear",
        align_corners=False,
    )
    return [*global_views, local, local.flip(-2)]


@dataclass(frozen=True, eq=False)
class DINOv2ObjectiveStep:
    student_routes: tuple[tuple[str, str], ...]
    teacher_routes: tuple[tuple[str, str], ...]
    student_view_ids: tuple[str, ...]
    teacher_view_ids: tuple[str, ...]
    masked_student_view_ids: tuple[str, ...]
    dino_pairs: tuple[tuple[str, str], ...]
    ibot_pairs: tuple[tuple[str, str], ...]
    koleo_view_ids: tuple[str, ...]
    masks: tuple[torch.Tensor, ...]
    student_global_cls: tuple[torch.Tensor, ...]
    student_cls_logits: tuple[torch.Tensor, ...]
    student_patch_logits: tuple[torch.Tensor, ...]
    teacher_cls_values: tuple[torch.Tensor, ...]
    teacher_cls_logits: tuple[torch.Tensor, ...]
    teacher_patch_logits: tuple[torch.Tensor, ...]
    patch_center_batch_mean: torch.Tensor
    dino_loss: torch.Tensor
    ibot_loss: torch.Tensor
    koleo_loss: torch.Tensor
    total_loss: torch.Tensor


def dinov2_objective_step(
    model: DINOv2Toy,
    student_views: list[torch.Tensor],
    masks: list[torch.Tensor],
) -> DINOv2ObjectiveStep:
    if not isinstance(model, DINOv2Toy):
        raise ValueError("model must be a DINOv2Toy")
    if not isinstance(student_views, (list, tuple)) or len(student_views) != 4:
        raise ValueError(
            "student_views must contain global-0, global-1, local-0, and local-1"
        )
    if not isinstance(masks, (list, tuple)) or len(masks) != 2:
        raise ValueError("masks must contain one mask for each global view")

    student_view_ids = ("global-0", "global-1", "local-0", "local-1")
    teacher_view_ids = ("global-0", "global-1")
    student_routes = (
        ("global-0", "masked student"),
        ("global-1", "masked student"),
        ("local-0", "clean student"),
        ("local-1", "clean student"),
    )
    teacher_routes = (
        ("global-0", "clean teacher"),
        ("global-1", "clean teacher"),
    )
    dino_pairs = tuple(
        (teacher_view_ids[teacher_index], student_view_ids[student_index])
        for teacher_index, student_index in _dino_pair_indices(4, 2)
    )
    ibot_pairs = tuple(zip(teacher_view_ids, teacher_view_ids))

    for global_view, mask in zip(student_views[:2], masks):
        with torch.no_grad():
            clean_patches = model.student.patch_tokens(global_view)
        _validate_patch_mask(
            mask,
            clean_patches.shape[0],
            clean_patches.shape[1],
            clean_patches.device,
        )

    student_global_cls = []
    student_cls_logits = []
    student_patch_logits = []
    teacher_cls_values = []
    teacher_cls_logits = []
    teacher_patch_logits = []
    for global_view, mask in zip(student_views[:2], masks):
        student_cls, teacher_cls, patch_logits, teacher_patch = model(
            global_view,
            mask,
        )
        student_global_cls.append(student_cls)
        student_cls_logits.append(model.student_cls_head(student_cls))
        student_patch_logits.append(patch_logits)
        teacher_cls_values.append(teacher_cls)
        teacher_cls_logits.append(model.teacher_cls_head(teacher_cls))
        teacher_patch_logits.append(teacher_patch)
    for local_view in student_views[2:]:
        _, student_cls = model.student.forward_features(local_view)
        student_cls_logits.append(model.student_cls_head(student_cls))

    patch_center_batch_mean = torch.cat(
        [
            teacher_view.detach()[mask]
            for teacher_view, mask in zip(teacher_patch_logits, masks)
        ],
        dim=0,
    ).mean(dim=0).reshape(1, 1, -1)
    next_cls_center = model.cls_center.detach().clone()
    next_patch_center = model.patch_center.detach().clone()
    dino_loss = dino_cross_view_loss(
        student_cls_logits,
        teacher_cls_logits,
        next_cls_center,
        student_temperature=model.student_temperature,
        teacher_temperature=model.teacher_temperature,
        center_momentum=model.center_momentum,
    )
    ibot_loss = ibot_masked_loss(
        student_patch_logits,
        teacher_patch_logits,
        masks,
        next_patch_center,
        student_temperature=model.student_temperature,
        teacher_temperature=model.teacher_temperature,
        center_momentum=model.center_momentum,
    )
    koleo = dinov2_koleo_loss(student_global_cls)
    total_loss = dino_loss + ibot_loss + 0.1 * koleo
    for loss in (dino_loss, ibot_loss, koleo, total_loss):
        if not bool(torch.isfinite(loss).item()):
            raise RuntimeError("DINOv2 objective losses must be finite")
    with torch.no_grad():
        model.cls_center.copy_(next_cls_center)
        model.patch_center.copy_(next_patch_center)

    return DINOv2ObjectiveStep(
        student_routes=student_routes,
        teacher_routes=teacher_routes,
        student_view_ids=student_view_ids,
        teacher_view_ids=teacher_view_ids,
        masked_student_view_ids=teacher_view_ids,
        dino_pairs=dino_pairs,
        ibot_pairs=ibot_pairs,
        koleo_view_ids=teacher_view_ids,
        masks=tuple(masks),
        student_global_cls=tuple(student_global_cls),
        student_cls_logits=tuple(student_cls_logits),
        student_patch_logits=tuple(student_patch_logits),
        teacher_cls_values=tuple(teacher_cls_values),
        teacher_cls_logits=tuple(teacher_cls_logits),
        teacher_patch_logits=tuple(teacher_patch_logits),
        patch_center_batch_mean=patch_center_batch_mean.detach(),
        dino_loss=dino_loss,
        ibot_loss=ibot_loss,
        koleo_loss=koleo,
        total_loss=total_loss,
    )


def _assert_probability_targets(targets: list[torch.Tensor]) -> None:
    for target in targets:
        assert not target.requires_grad
        assert torch.isfinite(target).all()
        assert torch.allclose(
            target.sum(dim=-1),
            torch.ones_like(target[..., 0]),
        )


def _assert_gradient(loss: torch.Tensor, parameter: torch.Tensor) -> None:
    gradient = torch.autograd.grad(
        loss,
        parameter,
        retain_graph=True,
    )[0]
    assert torch.isfinite(gradient).all()
    assert torch.count_nonzero(gradient) > 0


def run_dino_demo(config: ToyConfig) -> ShapeReport:
    model = DINOStudentTeacherToy(config).cpu()
    images = synthetic_images(config)
    student_views = _ssl_student_views(images, config)
    teacher_views = student_views[:2]
    assert len(student_views) == 4
    assert len(teacher_views) == 2
    assert _dino_pair_indices(len(student_views), len(teacher_views)) == [
        (0, 1),
        (0, 2),
        (0, 3),
        (1, 0),
        (1, 2),
        (1, 3),
    ]

    student_logits = [model.student_logits(view) for view in student_views]
    teacher_logits = [model.teacher_logits(view) for view in teacher_views]
    assert all(not value.requires_grad for value in teacher_logits)
    old_center = model.center.detach().clone()
    targets = [
        (
            (value.detach() - old_center) / model.teacher_temperature
        ).softmax(dim=-1)
        for value in teacher_logits
    ]
    _assert_probability_targets(targets)
    loss = dino_cross_view_loss(
        student_logits,
        teacher_logits,
        model.center,
        student_temperature=model.student_temperature,
        teacher_temperature=model.teacher_temperature,
        center_momentum=model.center_momentum,
    )
    assert torch.isfinite(loss)
    expected_center = model.center_momentum * old_center + (
        1 - model.center_momentum
    ) * torch.cat(teacher_logits).mean(dim=0, keepdim=True)
    assert torch.allclose(model.center, expected_center)

    optimizer = torch.optim.SGD(
        [parameter for parameter in model.parameters() if parameter.requires_grad],
        lr=0.01,
    )
    loss.backward()
    for parameter in (
        model.student.patch_embed.proj.weight,
        model.student_head.weight,
    ):
        assert parameter.grad is not None
        assert torch.isfinite(parameter.grad).all()
        assert torch.count_nonzero(parameter.grad) > 0
    assert all(parameter.grad is None for parameter in model.teacher.parameters())
    assert all(
        parameter.grad is None for parameter in model.teacher_head.parameters()
    )

    teacher_before = [
        parameter.detach().clone() for parameter in model.teacher.parameters()
    ]
    teacher_head_before = [
        parameter.detach().clone() for parameter in model.teacher_head.parameters()
    ]
    optimizer.step()
    student_after = [
        parameter.detach().clone() for parameter in model.student.parameters()
    ]
    student_head_after = [
        parameter.detach().clone() for parameter in model.student_head.parameters()
    ]
    ema_momentum = 0.9
    ema_update(model.teacher, model.student, ema_momentum)
    ema_update(model.teacher_head, model.student_head, ema_momentum)
    for actual, old, post_step in zip(
        model.teacher.parameters(),
        teacher_before,
        student_after,
    ):
        assert torch.allclose(
            actual,
            ema_momentum * old + (1 - ema_momentum) * post_step,
        )
    for actual, old, post_step in zip(
        model.teacher_head.parameters(),
        teacher_head_before,
        student_head_after,
    ):
        assert torch.allclose(
            actual,
            ema_momentum * old + (1 - ema_momentum) * post_step,
        )
    assert any(
        not torch.equal(actual, old)
        for actual, old in zip(
            model.teacher.parameters(),
            teacher_before,
        )
    )
    return {
        "student_views": (len(student_views),),
        "teacher_views": (len(teacher_views),),
        "student_logits": tuple(student_logits[0].shape),
    }


def run_dinov2_demo(config: ToyConfig) -> ShapeReport:
    model = DINOv2Toy(config).cpu()
    images = synthetic_images(config)
    student_views = _ssl_student_views(images, config)
    global_views = student_views[:2]
    local_views = student_views[2:]
    assert len(global_views) == 2
    assert len(local_views) == 2

    masks = []
    _, expected_masked = mask_partition_sizes(config.patch_tokens, 0.4)
    for view in global_views:
        with torch.no_grad():
            clean_patches = model.student.patch_tokens(view)
            _, _, mask = random_mask(clean_patches, 0.4)
        assert torch.equal(
            mask.sum(dim=1),
            torch.full(
                (config.batch_size,),
                expected_masked,
                device=mask.device,
            ),
        )
        masks.append(mask)

    old_cls_center = model.cls_center.detach().clone()
    old_patch_center = model.patch_center.detach().clone()
    step = dinov2_objective_step(model, student_views, masks=masks)
    assert step.student_routes == (
        ("global-0", "masked student"),
        ("global-1", "masked student"),
        ("local-0", "clean student"),
        ("local-1", "clean student"),
    )
    assert step.teacher_routes == (
        ("global-0", "clean teacher"),
        ("global-1", "clean teacher"),
    )
    assert step.dino_pairs == (
        ("global-0", "global-1"),
        ("global-0", "local-0"),
        ("global-0", "local-1"),
        ("global-1", "global-0"),
        ("global-1", "local-0"),
        ("global-1", "local-1"),
    )
    assert step.ibot_pairs == (
        ("global-0", "global-0"),
        ("global-1", "global-1"),
    )
    assert step.koleo_view_ids == ("global-0", "global-1")
    assert len(step.student_cls_logits) == 4
    assert len(step.teacher_cls_logits) == 2
    assert all(not value.requires_grad for value in step.teacher_cls_logits)
    assert all(not value.requires_grad for value in step.teacher_patch_logits)

    cls_targets = [
        (
            (value.detach() - old_cls_center) / model.teacher_temperature
        ).softmax(dim=-1)
        for value in step.teacher_cls_logits
    ]
    patch_targets = [
        (
            (value.detach() - old_patch_center) / model.teacher_temperature
        ).softmax(dim=-1)
        for value in step.teacher_patch_logits
    ]
    _assert_probability_targets(cls_targets)
    _assert_probability_targets(patch_targets)

    for loss in (step.dino_loss, step.ibot_loss, step.koleo_loss):
        assert torch.isfinite(loss)

    expected_cls_center = model.center_momentum * old_cls_center + (
        1 - model.center_momentum
    ) * torch.cat(step.teacher_cls_logits).mean(dim=0, keepdim=True)
    expected_patch_center = model.center_momentum * old_patch_center + (
        1 - model.center_momentum
    ) * step.patch_center_batch_mean
    assert torch.allclose(model.cls_center, expected_cls_center)
    assert torch.allclose(model.patch_center, expected_patch_center)

    _assert_gradient(step.dino_loss, model.student_cls_head.weight)
    _assert_gradient(step.ibot_loss, model.student_patch_head.weight)
    _assert_gradient(step.koleo_loss, model.student.patch_embed.proj.weight)
    optimizer = torch.optim.SGD(
        [parameter for parameter in model.parameters() if parameter.requires_grad],
        lr=0.01,
    )
    step.total_loss.backward()
    for parameter in (
        model.student.patch_embed.proj.weight,
        model.student_cls_head.weight,
        model.student_patch_head.weight,
        model.mask_token,
    ):
        assert parameter.grad is not None
        assert torch.isfinite(parameter.grad).all()
        assert torch.count_nonzero(parameter.grad) > 0
    assert all(parameter.grad is None for parameter in model.teacher.parameters())
    assert all(
        parameter.grad is None for parameter in model.teacher_cls_head.parameters()
    )
    assert all(
        parameter.grad is None
        for parameter in model.teacher_patch_head.parameters()
    )

    teacher_before = [
        parameter.detach().clone() for parameter in model.teacher.parameters()
    ]
    teacher_cls_head_before = [
        parameter.detach().clone()
        for parameter in model.teacher_cls_head.parameters()
    ]
    teacher_patch_head_before = [
        parameter.detach().clone()
        for parameter in model.teacher_patch_head.parameters()
    ]
    optimizer.step()
    student_after = [
        parameter.detach().clone() for parameter in model.student.parameters()
    ]
    student_cls_head_after = [
        parameter.detach().clone()
        for parameter in model.student_cls_head.parameters()
    ]
    student_patch_head_after = [
        parameter.detach().clone()
        for parameter in model.student_patch_head.parameters()
    ]
    ema_momentum = 0.9
    ema_update(model.teacher, model.student, ema_momentum)
    ema_update(model.teacher_cls_head, model.student_cls_head, ema_momentum)
    ema_update(model.teacher_patch_head, model.student_patch_head, ema_momentum)
    for actual, old, post_step in zip(
        model.teacher.parameters(), teacher_before, student_after
    ):
        assert torch.allclose(
            actual,
            ema_momentum * old + (1 - ema_momentum) * post_step,
        )
    for actual, old, post_step in zip(
        model.teacher_cls_head.parameters(),
        teacher_cls_head_before,
        student_cls_head_after,
    ):
        assert torch.allclose(
            actual,
            ema_momentum * old + (1 - ema_momentum) * post_step,
        )
    for actual, old, post_step in zip(
        model.teacher_patch_head.parameters(),
        teacher_patch_head_before,
        student_patch_head_after,
    ):
        assert torch.allclose(
            actual,
            ema_momentum * old + (1 - ema_momentum) * post_step,
        )
    return {
        "student_cls": tuple(step.student_global_cls[0].shape),
        "teacher_cls": tuple(step.teacher_cls_values[0].shape),
        "patch_logits": tuple(step.student_patch_logits[0].shape),
        "mask": tuple(step.masks[0].shape),
    }


DEMOS["dino"] = run_dino_demo
DEMOS["dinov2"] = run_dinov2_demo


class RegisteredViT(TinyViT):
    def __init__(self, config: ToyConfig, register_count: int = 4):
        if isinstance(register_count, bool) or not isinstance(register_count, int):
            raise TypeError("register_count must be a non-negative integer")
        if register_count < 0:
            raise ValueError("register_count must be a non-negative integer")
        super().__init__(config)
        self.register_count = register_count
        self.register_tokens = nn.Parameter(
            torch.empty(1, register_count, config.width)
        )
        if register_count:
            nn.init.normal_(self.register_tokens, std=0.02)

    def prepend_tokens(self, patches: Tensor) -> Tensor:
        batch, patch_count, width = patches.shape
        if patch_count != self.config.patch_tokens or width != self.config.width:
            raise ValueError("patches must match the registered ViT configuration")
        tokens = torch.cat(
            (
                self.cls_token.expand(batch, -1, -1),
                self.register_tokens.expand(batch, -1, -1),
                patches,
            ),
            dim=1,
        )
        assert tokens.shape == (
            batch,
            1 + self.register_count + self.config.patch_tokens,
            self.config.width,
        )
        return tokens

    def add_positions(self, tokens: Tensor) -> Tensor:
        expected_tokens = 1 + self.register_count + self.config.patch_tokens
        if tokens.shape[1:] != (expected_tokens, self.config.width):
            raise ValueError("tokens must use CLS | registers | patches order")
        cls = tokens[:, :1] + self.position[:, :1]
        registers = tokens[:, 1 : 1 + self.register_count]
        patches = tokens[:, 1 + self.register_count :] + self.position[:, 1:]
        return torch.cat((cls, registers, patches), dim=1)


def run_registers_demo(config: ToyConfig) -> ShapeReport:
    model = RegisteredViT(config).cpu()
    images = synthetic_images(config)
    tokens, cls = model.forward_features(images)
    registers = tokens[:, 1 : 1 + model.register_count]
    patches = tokens[:, 1 + model.register_count :]
    logits = model.head(cls)
    labels = torch.arange(config.batch_size) % config.classes
    loss = F.cross_entropy(logits, labels)
    if not torch.isfinite(loss):
        raise RuntimeError("register demo loss must be finite")
    loss.backward()

    assert RegisteredViT.forward_features is TinyViT.forward_features
    assert model.position.shape == (
        1,
        1 + config.patch_tokens,
        config.width,
    )
    assert tokens.shape == (
        config.batch_size,
        1 + model.register_count + config.patch_tokens,
        config.width,
    )
    assert registers.shape == (
        config.batch_size,
        model.register_count,
        config.width,
    )
    assert patches.shape == (
        config.batch_size,
        config.patch_tokens,
        config.width,
    )
    assert cls.shape == (config.batch_size, config.width)
    assert logits.shape == (config.batch_size, config.classes)
    for gradient in (
        model.patch_embed.proj.weight.grad,
        model.register_tokens.grad,
        model.blocks[0].attn.qkv.weight.grad,
    ):
        assert gradient is not None
        assert torch.isfinite(gradient).all()
        assert torch.count_nonzero(gradient) > 0
    return {
        "tokens": tuple(tokens.shape),
        "registers": tuple(registers.shape),
        "patches": tuple(patches.shape),
        "logits": tuple(logits.shape),
    }


DEMOS["registers"] = run_registers_demo


def _validate_clip_feature_pair(
    image_features: Tensor,
    text_features: Tensor,
) -> None:
    if not isinstance(image_features, Tensor) or not isinstance(
        text_features, Tensor
    ):
        raise TypeError("CLIP image_features and text_features must be tensors")
    if image_features.ndim != 2 or text_features.ndim != 2:
        raise ValueError("CLIP image/text features must be rank-2 tensors")
    if not torch.is_floating_point(image_features) or not torch.is_floating_point(
        text_features
    ):
        raise TypeError("CLIP image/text features must use floating dtypes")
    if image_features.shape != text_features.shape:
        raise ValueError("CLIP image/text feature shapes must match")
    if image_features.shape[0] < 1:
        raise ValueError("CLIP features must contain at least one example")
    if image_features.shape[1] < 1:
        raise ValueError("CLIP feature width must be non-empty")
    if image_features.dtype != text_features.dtype:
        raise TypeError("CLIP image/text features must use the same dtype")
    if image_features.device != text_features.device:
        raise ValueError("CLIP image/text features must use the same device")
    if not bool(torch.isfinite(image_features).all().item()) or not bool(
        torch.isfinite(text_features).all().item()
    ):
        raise ValueError("CLIP image/text features must be finite")


def _validate_clip_temperature(temperature: float) -> None:
    if isinstance(temperature, bool) or not isinstance(temperature, (int, float)):
        raise TypeError("CLIP temperature must be a non-boolean number")
    try:
        temperature_is_finite = math.isfinite(temperature)
    except OverflowError as error:
        raise ValueError("CLIP temperature must be positive and finite") from error
    if not temperature_is_finite or temperature <= 0:
        raise ValueError("CLIP temperature must be positive and finite")


def _clip_directional_loss(
    image_normalized: Tensor,
    text_normalized: Tensor,
    temperature: float,
) -> tuple[Tensor, Tensor]:
    logits = image_normalized @ text_normalized.transpose(0, 1) / temperature
    if not bool(torch.isfinite(logits).all().item()):
        raise ValueError(
            "CLIP temperature must keep normalized similarity logits finite"
        )
    targets = torch.arange(logits.shape[0], device=logits.device)
    image_to_text = F.cross_entropy(logits, targets)
    text_to_image = F.cross_entropy(logits.transpose(0, 1), targets)
    loss = 0.5 * (image_to_text + text_to_image)
    if not bool(torch.isfinite(loss).item()):
        raise ValueError("CLIP temperature must produce a finite loss")
    return loss, logits


def _clip_symmetric_loss_from_normalized(
    image_normalized: Tensor,
    text_normalized: Tensor,
    temperature: float = 0.07,
) -> tuple[Tensor, Tensor]:
    """Consume features that must already be L2-normalized exactly once."""

    _validate_clip_feature_pair(image_normalized, text_normalized)
    _validate_clip_temperature(temperature)
    for name, features in (
        ("image_normalized", image_normalized),
        ("text_normalized", text_normalized),
    ):
        norms = features.norm(dim=-1)
        if not torch.allclose(
            norms,
            torch.ones_like(norms),
            rtol=1e-5,
            atol=1e-6,
        ):
            raise ValueError(f"CLIP {name} must already be L2-normalized")
    return _clip_directional_loss(image_normalized, text_normalized, temperature)


def clip_symmetric_loss(
    image_features: Tensor,
    text_features: Tensor,
    temperature: float = 0.07,
) -> tuple[Tensor, Tensor]:
    """Defensively normalize arbitrary features for the symmetric CLIP loss.

    This toy fixes its default temperature at 0.07; reference CLIP learns a
    logit scale. The runnable demo uses the normalized-feature seam above so
    its explicit projection normalization is the only normalization pass.
    """

    _validate_clip_feature_pair(image_features, text_features)
    _validate_clip_temperature(temperature)
    return _clip_directional_loss(
        F.normalize(image_features, dim=-1),
        F.normalize(text_features, dim=-1),
        temperature,
    )


def run_clip_demo(config: ToyConfig) -> ShapeReport:
    if config.batch_size < 2:
        raise ValueError("CLIP demo needs at least two image-text pairs")
    vision = TinyViT(config).cpu()
    images = synthetic_images(config)
    tokens, image_features = vision.forward_features(images)
    text_tower = nn.Embedding(config.batch_size, config.width).cpu()
    text_ids = torch.arange(config.batch_size)
    text_features = text_tower(text_ids)
    image_projection = nn.Linear(config.width, config.width, bias=False).cpu()
    text_projection = nn.Linear(config.width, config.width, bias=False).cpu()
    assert image_projection is not text_projection
    assert image_projection.weight.data_ptr() != text_projection.weight.data_ptr()
    image_projected = image_projection(image_features)
    text_projected = text_projection(text_features)
    image_normalized = F.normalize(image_projected, dim=-1)
    text_normalized = F.normalize(text_projected, dim=-1)
    assert torch.allclose(
        image_normalized.norm(dim=-1),
        torch.ones(config.batch_size),
        atol=1e-6,
    )
    assert torch.allclose(
        text_normalized.norm(dim=-1),
        torch.ones(config.batch_size),
        atol=1e-6,
    )
    loss, logits = _clip_symmetric_loss_from_normalized(
        image_normalized,
        text_normalized,
    )
    if not bool(torch.isfinite(loss).item()):
        raise RuntimeError("CLIP demo loss must be finite")
    loss.backward()

    assert tokens.shape == (
        config.batch_size,
        config.patch_tokens + 1,
        config.width,
    )
    assert image_features.shape == (config.batch_size, config.width)
    assert text_features.shape == (config.batch_size, config.width)
    assert logits.shape == (config.batch_size, config.batch_size)
    for gradient in (
        vision.patch_embed.proj.weight.grad,
        text_tower.weight.grad,
        image_projection.weight.grad,
        text_projection.weight.grad,
    ):
        assert gradient is not None
        assert torch.isfinite(gradient).all()
        assert torch.count_nonzero(gradient) > 0
    return {
        "image_features": tuple(image_features.shape),
        "logits": tuple(logits.shape),
        "tokens": tuple(tokens.shape),
    }


DEMOS["clip"] = run_clip_demo


def siglip_pair_loss(
    image_features: Tensor,
    text_features: Tensor,
    bias: float = 0.0,
) -> tuple[Tensor, Tensor, Tensor]:
    """Compute SigLIP Algorithm 1 with a fixed unit logit scale.

    The paper learns an inverse temperature and bias. This small teaching seam
    keeps the scale at one and exposes only the additive bias so the pairwise
    sigmoid objective stays visible.
    """
    if not isinstance(image_features, Tensor) or not isinstance(
        text_features, Tensor
    ):
        raise TypeError("SigLIP image_features and text_features must be tensors")
    if image_features.ndim != 2 or text_features.ndim != 2:
        raise ValueError("SigLIP image/text features must be rank-2 tensors")
    if not torch.is_floating_point(image_features) or not torch.is_floating_point(
        text_features
    ):
        raise TypeError("SigLIP image/text features must use floating dtypes")
    if image_features.shape != text_features.shape:
        raise ValueError("SigLIP image/text feature shapes must match")
    if image_features.shape[0] < 1:
        raise ValueError("SigLIP features must contain at least one example")
    if image_features.shape[1] < 1:
        raise ValueError("SigLIP feature width must be non-empty")
    if image_features.dtype != text_features.dtype:
        raise TypeError("SigLIP image/text features must use the same dtype")
    if image_features.device != text_features.device:
        raise ValueError("SigLIP image/text features must use the same device")
    if not bool(torch.isfinite(image_features).all().item()) or not bool(
        torch.isfinite(text_features).all().item()
    ):
        raise ValueError("SigLIP image/text features must be finite")
    if isinstance(bias, bool) or not isinstance(bias, (int, float)):
        raise TypeError("SigLIP bias must be a non-boolean number")
    try:
        finite_bias = math.isfinite(bias)
    except OverflowError as error:
        raise ValueError(
            "SigLIP bias must be finite and representable"
        ) from error
    if not finite_bias:
        raise ValueError("SigLIP bias must be finite")

    image_normalized = F.normalize(image_features, dim=-1)
    text_normalized = F.normalize(text_features, dim=-1)
    try:
        bias_tensor = torch.tensor(
            bias,
            dtype=image_features.dtype,
            device=image_features.device,
        )
    except (OverflowError, RuntimeError) as error:
        raise ValueError(
            "SigLIP bias must be representable in the feature dtype"
        ) from error
    if not bool(torch.isfinite(bias_tensor).item()):
        raise ValueError("SigLIP bias must be representable and finite")

    logits = image_normalized @ text_normalized.transpose(0, 1) + bias_tensor
    if not bool(torch.isfinite(logits).all().item()):
        raise ValueError("SigLIP bias must keep pair logits finite")
    labels = 2.0 * torch.eye(
        logits.shape[0],
        dtype=logits.dtype,
        device=logits.device,
    ) - 1.0
    pair_losses = -F.logsigmoid(labels * logits)
    loss = pair_losses.sum() / logits.shape[0]
    if not bool(torch.isfinite(loss).item()):
        raise ValueError(
            "SigLIP bias must keep the batch-normalized pair loss finite"
        )
    return loss, logits, labels


def run_siglip_demo(config: ToyConfig) -> ShapeReport:
    if config.batch_size < 2:
        raise ValueError("SigLIP demo needs at least two image-text pairs")
    vision = TinyViT(config).cpu()
    images = synthetic_images(config)
    tokens, image_features = vision.forward_features(images)
    text_tower = nn.Embedding(config.batch_size, config.width).cpu()
    text_ids = torch.arange(config.batch_size)
    text_features = text_tower(text_ids)
    loss, logits, labels = siglip_pair_loss(image_features, text_features)
    if not bool(torch.isfinite(loss).item()):
        raise RuntimeError("SigLIP demo loss must be finite")
    loss.backward()

    assert tokens.shape == (
        config.batch_size,
        config.patch_tokens + 1,
        config.width,
    )
    assert image_features.shape == (config.batch_size, config.width)
    assert text_features.shape == (config.batch_size, config.width)
    assert logits.shape == (config.batch_size, config.batch_size)
    assert labels.shape == logits.shape
    assert torch.equal(
        labels.diagonal(),
        torch.ones(config.batch_size, dtype=labels.dtype),
    )
    off_diagonal = ~torch.eye(config.batch_size, dtype=torch.bool)
    assert torch.equal(
        labels[off_diagonal],
        -torch.ones(config.batch_size**2 - config.batch_size, dtype=labels.dtype),
    )
    for gradient in (
        vision.patch_embed.proj.weight.grad,
        text_tower.weight.grad,
    ):
        assert gradient is not None
        assert torch.isfinite(gradient).all()
        assert torch.count_nonzero(gradient) > 0
    return {
        "labels": tuple(labels.shape),
        "logits": tuple(logits.shape),
    }


DEMOS["siglip"] = run_siglip_demo


def vitdet_feature_pyramid(
    patch_tokens: Tensor,
    grid_size: int,
) -> dict[str, Tensor]:
    if not isinstance(patch_tokens, Tensor):
        raise TypeError("ViTDet patch_tokens must be a tensor")
    if patch_tokens.ndim != 3:
        raise ValueError("ViTDet patch_tokens must be a rank-3 tensor")
    if not torch.is_floating_point(patch_tokens):
        raise TypeError("ViTDet patch_tokens must use a floating dtype")
    if not bool(torch.isfinite(patch_tokens).all().item()):
        raise ValueError("ViTDet patch_tokens must be finite")
    if isinstance(grid_size, bool) or not isinstance(grid_size, int) or grid_size <= 0:
        raise ValueError("ViTDet grid_size must be a positive integer")
    batch, tokens, width = patch_tokens.shape
    if batch == 0 or width == 0:
        raise ValueError("ViTDet patch_tokens must have non-empty batch and width")
    if tokens != grid_size**2:
        raise ValueError("ViTDet patch token count must match grid_size²")
    if grid_size % 2 != 0:
        raise ValueError("ViTDet grid_size must be even for the low branch")

    feature = patch_tokens.transpose(1, 2).reshape(batch, width, grid_size, grid_size)
    high = F.interpolate(feature, scale_factor=2.0, mode="nearest")
    low = F.avg_pool2d(feature, kernel_size=2, stride=2)
    return {"high": high, "base": feature, "low": low}


def run_vitdet_demo(config: ToyConfig) -> ShapeReport:
    model = TinyViT(config).cpu()
    images = synthetic_images(config)
    tokens, _ = model.forward_features(images)
    pyramid = vitdet_feature_pyramid(tokens[:, 1:], config.grid_size)
    assert pyramid["high"].shape == (
        config.batch_size,
        config.width,
        config.grid_size * 2,
        config.grid_size * 2,
    )
    assert pyramid["base"].shape == (
        config.batch_size,
        config.width,
        config.grid_size,
        config.grid_size,
    )
    assert pyramid["low"].shape == (
        config.batch_size,
        config.width,
        config.grid_size // 2,
        config.grid_size // 2,
    )
    loss = sum(value.square().mean() for value in pyramid.values())
    assert bool(torch.isfinite(loss).item())
    loss.backward()
    gradient = model.patch_embed.proj.weight.grad
    assert gradient is not None
    assert bool(torch.isfinite(gradient).all().item())
    assert torch.count_nonzero(gradient) > 0
    return {name: tuple(value.shape) for name, value in pyramid.items()}


DEMOS["vitdet"] = run_vitdet_demo


class ToySAM(nn.Module):
    def __init__(self, config: ToyConfig) -> None:
        super().__init__()
        self.vision = TinyViT(config)
        self.prompt_projection = nn.Linear(2, config.width)
        self.mask_projection = nn.Linear(config.width, config.patch_tokens)
        self.image_encode_calls = 0

    def encode_image(self, images: Tensor) -> Tensor:
        self.image_encode_calls += 1
        tokens, _ = self.vision.forward_features(images)
        return tokens[:, 1:]

    def decode_prompt(self, image_tokens: Tensor, points: Tensor) -> Tensor:
        if not isinstance(image_tokens, Tensor):
            raise TypeError("image_tokens must be a tensor")
        if image_tokens.ndim != 3:
            raise ValueError("image_tokens must be a rank-3 tensor")
        if not torch.is_floating_point(image_tokens):
            raise TypeError("image_tokens must use a floating dtype")
        if any(dimension == 0 for dimension in image_tokens.shape):
            raise ValueError("image_tokens must have non-empty dimensions")
        parameter = self.prompt_projection.weight
        if image_tokens.shape[2] != parameter.shape[0]:
            raise ValueError("image_tokens width must match the model")
        if image_tokens.dtype != parameter.dtype:
            raise TypeError("image_tokens dtype must match the model")
        if image_tokens.device != parameter.device:
            raise ValueError("image_tokens device must match the model")
        if not bool(torch.isfinite(image_tokens).all().item()):
            raise ValueError("image_tokens must be finite")
        if not isinstance(points, Tensor):
            raise TypeError("points must be a tensor")
        if points.ndim != 2 or points.shape[1] != 2:
            raise ValueError("points shape must be [B, 2]")
        if any(dimension == 0 for dimension in points.shape):
            raise ValueError("points must have non-empty dimensions")
        if not torch.is_floating_point(points):
            raise TypeError("points dtype must be floating")
        if image_tokens.shape[0] != points.shape[0]:
            raise ValueError("image_tokens and points batch dimensions must match")
        if points.dtype != parameter.dtype:
            raise TypeError("points dtype must match the model")
        if points.device != parameter.device:
            raise ValueError("points device must match the model")
        if not bool(torch.isfinite(points).all().item()):
            raise ValueError("points finite values are required")

        prompt = self.prompt_projection(points).unsqueeze(1)
        pooled = image_tokens.mean(dim=1, keepdim=True) + prompt
        return self.mask_projection(pooled.squeeze(1))


def run_sam_demo(config: ToyConfig) -> ShapeReport:
    model = ToySAM(config).cpu()
    images = synthetic_images(config)
    image_tokens = model.encode_image(images)
    points = [
        torch.tensor(
            [[float(index), float(index + 1)] for _ in range(config.batch_size)]
        )
        for index in range(3)
    ]
    masks = [model.decode_prompt(image_tokens, point) for point in points]
    assert model.image_encode_calls == 1
    assert all(mask.shape == (config.batch_size, config.patch_tokens) for mask in masks)
    assert not torch.allclose(masks[0], masks[1])
    loss = sum(mask.square().mean() for mask in masks)
    assert bool(torch.isfinite(loss).item())
    loss.backward()
    for gradient in (
        model.vision.patch_embed.proj.weight.grad,
        model.prompt_projection.weight.grad,
        model.mask_projection.weight.grad,
    ):
        assert gradient is not None
        assert bool(torch.isfinite(gradient).all().item())
        assert torch.count_nonzero(gradient) > 0
    return {
        "image_tokens": tuple(image_tokens.shape),
        "masks": (len(masks), *tuple(masks[0].shape)),
    }


DEMOS["sam"] = run_sam_demo


def run_blip2_demo(config: ToyConfig) -> ShapeReport:
    query_tokens = BLIP2_TOY_QUERY_TOKENS
    assert query_tokens == 8
    vision = TinyViT(config).cpu()
    vision.requires_grad_(False)
    vision.eval()
    images = synthetic_images(config)
    with torch.no_grad():
        source_tokens, _ = vision.forward_features(images)
    resampler = LearnedQueryResampler(
        config.width,
        config.heads,
        query_tokens=query_tokens,
    ).cpu()
    output, weights = resampler(source_tokens)
    assert source_tokens.shape == (
        config.batch_size,
        config.patch_tokens + 1,
        config.width,
    )
    assert output.shape == (config.batch_size, query_tokens, config.width)
    assert weights.shape == (
        config.batch_size,
        query_tokens,
        config.patch_tokens + 1,
    )
    assert bool(torch.isfinite(weights).all().item())
    assert bool((weights >= 0).all().item())
    assert torch.allclose(
        weights.sum(dim=-1),
        torch.ones(config.batch_size, query_tokens, dtype=weights.dtype),
    )
    loss = output.square().mean() + weights.square().mean()
    assert bool(torch.isfinite(loss).item())
    loss.backward()
    assert all(parameter.grad is None for parameter in vision.parameters())
    for gradient in (
        resampler.queries.grad,
        resampler.cross_attention.in_proj_weight.grad,
        resampler.cross_attention.out_proj.weight.grad,
    ):
        assert gradient is not None
        assert bool(torch.isfinite(gradient).all().item())
        assert torch.count_nonzero(gradient) > 0
    return {
        "attention": tuple(weights.shape),
        "query_tokens": tuple(output.shape),
        "source_tokens": tuple(source_tokens.shape),
    }


DEMOS["blip2"] = run_blip2_demo


LLAVA_TOY_TEXT_TOKENS = 6
LLAVA_TOY_LANGUAGE_WIDTH = 80
LLAVA_TOY_DECODE_APPEND_STEPS = 2
LLAVA_TOY_LAYERS = 2
LLAVA_TOY_KV_HEADS = 2
LLAVA_TOY_HEAD_WIDTH = 8


def _llava_rank3_floating_finite(value: object, name: str) -> Tensor:
    if not isinstance(value, Tensor):
        raise TypeError(f"{name} must be a tensor")
    if value.ndim != 3:
        raise ValueError(f"{name} must be a rank-3 tensor")
    if not torch.is_floating_point(value):
        raise TypeError(f"{name} must use a floating dtype")
    if any(dimension == 0 for dimension in value.shape):
        raise ValueError(f"{name} must have non-empty batch, sequence, and width")
    if value.device.type != "meta" and not bool(torch.isfinite(value).all().item()):
        raise ValueError(f"{name} must be finite")
    return value


def llava_visual_prefix(
    image_tokens: Tensor,
    text_embeddings: Tensor,
    projector: nn.Module,
) -> Tensor:
    """Use a teaching-only image-first concatenation layout.

    Reference prompts may place the image placeholder before or after the question.
    Implementations insert image features at that placeholder. This helper
    deliberately demonstrates only the image-first case.
    """

    image_tokens = _llava_rank3_floating_finite(image_tokens, "image_tokens")
    text_embeddings = _llava_rank3_floating_finite(
        text_embeddings, "text_embeddings"
    )
    if image_tokens.shape[0] != text_embeddings.shape[0]:
        raise ValueError("LLaVA image/text batches must match")
    if image_tokens.dtype != text_embeddings.dtype:
        raise TypeError("LLaVA image/text tensors must use the same dtype")
    if image_tokens.device != text_embeddings.device:
        raise ValueError("LLaVA image/text tensors must use the same device")
    if not isinstance(projector, nn.Module):
        raise TypeError("LLaVA projector must be an nn.Module")
    parameter = next(projector.parameters(), None)
    if parameter is not None:
        if image_tokens.dtype != parameter.dtype:
            raise TypeError("LLaVA projector dtype must match image tokens")
        if image_tokens.device != parameter.device:
            raise ValueError("LLaVA projector device must match image tokens")
    projected = projector(image_tokens)
    if not isinstance(projected, Tensor) or projected.ndim != 3:
        raise ValueError("LLaVA projector output must be a rank-3 tensor")
    if projected.shape[:2] != image_tokens.shape[:2]:
        raise ValueError("LLaVA projector output batch and sequence must match images")
    if projected.shape[-1] != text_embeddings.shape[-1]:
        raise ValueError("LLaVA projector must emit the language width")
    if projected.dtype != text_embeddings.dtype:
        raise TypeError("LLaVA projector output dtype must match text embeddings")
    if projected.device != text_embeddings.device:
        raise ValueError("LLaVA projector output device must match text embeddings")
    if not bool(torch.isfinite(projected).all().item()):
        raise ValueError("LLaVA projector output must be finite")
    return torch.cat((projected, text_embeddings), dim=1)


class ToyDecoderKVCache(nn.Module):
    def __init__(
        self,
        embedding_width: int,
        layers: int,
        kv_heads: int,
        head_width: int,
    ) -> None:
        super().__init__()
        for name, value in (
            ("embedding_width", embedding_width),
            ("layers", layers),
            ("kv_heads", kv_heads),
            ("head_width", head_width),
        ):
            if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
                raise ValueError(f"{name} must be a positive integer")
        self.embedding_width = embedding_width
        self.layers = layers
        self.kv_heads = kv_heads
        self.head_width = head_width
        projection_width = layers * kv_heads * head_width
        self.key_projection = nn.Linear(embedding_width, projection_width)
        self.value_projection = nn.Linear(embedding_width, projection_width)
        self._key: Tensor | None = None
        self._value: Tensor | None = None
        self.prefill_calls = 0
        self.append_calls = 0
        self.cache_length_history: list[int] = []
        self.projection_token_counts: list[int] = []

    def _validate_embeddings(
        self,
        embeddings: object,
        *,
        expected_batch: int | None = None,
        one_token: bool = False,
    ) -> Tensor:
        value = _llava_rank3_floating_finite(embeddings, "decoder embeddings")
        parameter = self.key_projection.weight
        if value.shape[-1] != self.embedding_width:
            raise ValueError("decoder embeddings width must match the cache")
        if value.dtype != parameter.dtype:
            raise TypeError("decoder embeddings dtype must match the cache")
        if value.device != parameter.device:
            raise ValueError("decoder embeddings device must match the cache")
        if expected_batch is not None and value.shape[0] != expected_batch:
            raise ValueError("decoder embeddings batch must match the prefetched cache")
        if one_token and value.shape[1] != 1:
            raise ValueError("decoder cache append must project exactly one token")
        return value

    def _project(self, embeddings: Tensor) -> tuple[Tensor, Tensor]:
        batch, tokens, _ = embeddings.shape
        shape = (batch, tokens, self.layers, self.kv_heads, self.head_width)
        key = self.key_projection(embeddings).reshape(shape).permute(0, 2, 3, 1, 4)
        value = (
            self.value_projection(embeddings).reshape(shape).permute(0, 2, 3, 1, 4)
        )
        assert key.shape == value.shape == (
            batch,
            self.layers,
            self.kv_heads,
            tokens,
            self.head_width,
        )
        self.projection_token_counts.append(tokens)
        return key, value

    def prefill(self, prefix_embeddings: Tensor) -> tuple[Tensor, Tensor]:
        if self._key is not None or self._value is not None:
            raise RuntimeError("decoder cache prefill runs exactly once")
        prefix_embeddings = self._validate_embeddings(prefix_embeddings)
        self._key, self._value = self._project(prefix_embeddings)
        self.prefill_calls += 1
        self.cache_length_history.append(self._key.shape[3])
        return self._key, self._value

    def append(self, generated_embedding: Tensor) -> tuple[Tensor, Tensor]:
        if self._key is None or self._value is None:
            raise RuntimeError("decoder cache append needs prefill")
        generated_embedding = self._validate_embeddings(
            generated_embedding,
            expected_batch=self._key.shape[0],
            one_token=True,
        )
        next_key, next_value = self._project(generated_embedding)
        self._key = torch.cat((self._key, next_key), dim=3)
        self._value = torch.cat((self._value, next_value), dim=3)
        self.append_calls += 1
        self.cache_length_history.append(self._key.shape[3])
        return self._key, self._value


def run_llava_demo(config: ToyConfig) -> ShapeReport:
    text_tokens = LLAVA_TOY_TEXT_TOKENS
    language_width = LLAVA_TOY_LANGUAGE_WIDTH
    decode_append_steps = LLAVA_TOY_DECODE_APPEND_STEPS
    layers = LLAVA_TOY_LAYERS
    kv_heads = LLAVA_TOY_KV_HEADS
    head_width = LLAVA_TOY_HEAD_WIDTH
    assert (
        text_tokens,
        language_width,
        decode_append_steps,
        layers,
        kv_heads,
        head_width,
    ) == (6, 80, 2, 2, 2, 8)
    vision = TinyViT(config).cpu()
    images = synthetic_images(config).requires_grad_()
    image_tokens, _ = vision.forward_features(images)
    projector = nn.Linear(config.width, language_width).cpu()
    text_embeddings = torch.randn(
        config.batch_size,
        text_tokens,
        language_width,
        requires_grad=True,
    )
    prefix = llava_visual_prefix(image_tokens[:, 1:], text_embeddings, projector)
    cache = ToyDecoderKVCache(
        language_width,
        layers=layers,
        kv_heads=kv_heads,
        head_width=head_width,
    ).cpu()
    prefill_key, prefill_value = cache.prefill(prefix)
    prefix_key = prefill_key.detach().clone()
    prefix_value = prefill_value.detach().clone()
    # These are explicit one-token append forwards, not a count of generated
    # outputs: in normal generation, prefill already produces the first output.
    append_embeddings = [
        torch.randn(config.batch_size, 1, language_width, requires_grad=True)
        for _ in range(decode_append_steps)
    ]
    post_decode_key, post_decode_value = prefill_key, prefill_value
    for embedding in append_embeddings:
        post_decode_key, post_decode_value = cache.append(embedding)

    assert image_tokens.shape == (
        config.batch_size,
        config.patch_tokens + 1,
        config.width,
    )
    assert prefix.shape == (
        config.batch_size,
        config.patch_tokens + text_tokens,
        language_width,
    )
    assert prefill_key.shape == prefill_value.shape == (
        config.batch_size,
        layers,
        kv_heads,
        config.patch_tokens + text_tokens,
        head_width,
    )
    assert post_decode_key.shape == post_decode_value.shape == (
        config.batch_size,
        layers,
        kv_heads,
        config.patch_tokens + text_tokens + decode_append_steps,
        head_width,
    )
    assert torch.equal(post_decode_key[:, :, :, : prefix.shape[1]], prefix_key)
    assert torch.equal(post_decode_value[:, :, :, : prefix.shape[1]], prefix_value)
    assert post_decode_key.data_ptr() != post_decode_value.data_ptr()
    assert not torch.allclose(post_decode_key, post_decode_value)
    assert cache.prefill_calls == 1
    assert cache.append_calls == decode_append_steps
    assert cache.cache_length_history == [70, 71, 72]
    assert cache.projection_token_counts == [70, 1, 1]

    loss = (
        prefix.square().mean()
        + prefill_key.square().mean()
        + prefill_value.square().mean()
        + post_decode_key.square().mean()
        + post_decode_value.square().mean()
    )
    assert bool(torch.isfinite(loss).item())
    loss.backward()
    for gradient in (
        images.grad,
        vision.patch_embed.proj.weight.grad,
        text_embeddings.grad,
        projector.weight.grad,
        cache.key_projection.weight.grad,
        cache.value_projection.weight.grad,
        *(embedding.grad for embedding in append_embeddings),
    ):
        assert gradient is not None
        assert bool(torch.isfinite(gradient).all().item())
        assert torch.count_nonzero(gradient) > 0
    return {
        "combined_prefix": tuple(prefix.shape),
        "image_tokens": tuple(image_tokens.shape),
        "post_decode_key": tuple(post_decode_key.shape),
        "post_decode_value": tuple(post_decode_value.shape),
        "prefill_key": tuple(prefill_key.shape),
        "prefill_value": tuple(prefill_value.shape),
        "visual_prefix": tuple(prefix[:, : config.patch_tokens].shape),
    }


DEMOS["llava"] = run_llava_demo


def _efficient_positive_integer(value: object, name: str) -> int:
    if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
        raise ValueError(f"{name} must be a positive integer")
    return value


class ToyMobileViTBlock(nn.Module):
    def __init__(
        self,
        width: int,
        heads: int,
        *,
        transformer_width: int | None = None,
        patch_height: int = 2,
        patch_width: int = 2,
    ) -> None:
        super().__init__()
        self.width = _efficient_positive_integer(width, "width")
        heads = _efficient_positive_integer(heads, "heads")
        self.transformer_width = (
            self.width
            if transformer_width is None
            else _efficient_positive_integer(transformer_width, "transformer_width")
        )
        self.patch_height = _efficient_positive_integer(patch_height, "patch_height")
        self.patch_width = _efficient_positive_integer(patch_width, "patch_width")
        if self.transformer_width % heads != 0:
            raise ValueError("transformer_width must be divisible by heads")
        self.local = nn.Sequential(
            nn.Conv2d(self.width, self.width, 3, padding=1),
            nn.SiLU(),
        )
        self.projection = nn.Conv2d(self.width, self.transformer_width, 1)
        self.transformer = nn.TransformerEncoderLayer(
            d_model=self.transformer_width,
            nhead=heads,
            dim_feedforward=2 * self.transformer_width,
            dropout=0.0,
            batch_first=True,
            norm_first=True,
        )
        self.post_fold_projection = nn.Conv2d(
            self.transformer_width,
            self.width,
            1,
        )
        self.fusion = nn.Conv2d(2 * self.width, self.width, 3, padding=1)

    def _validate_features(
        self,
        features: object,
        *,
        width: int | None = None,
    ) -> Tensor:
        expected_width = self.width if width is None else width
        if not isinstance(features, Tensor):
            raise TypeError("MobileViT features must be a tensor")
        if features.ndim != 4:
            raise ValueError("MobileViT features must be rank-4")
        if not torch.is_floating_point(features):
            raise TypeError("MobileViT features must use a floating dtype")
        if any(dimension == 0 for dimension in features.shape):
            raise ValueError("MobileViT features must have non-empty dimensions")
        if features.shape[1] != expected_width:
            raise ValueError("MobileViT features width must match the block")
        if features.shape[2] % self.patch_height != 0:
            raise ValueError("MobileViT feature height must be divisible by patch_height")
        if features.shape[3] % self.patch_width != 0:
            raise ValueError("MobileViT feature width must be divisible by patch_width")
        parameter = self.projection.weight
        if features.device != parameter.device:
            raise ValueError("MobileViT features device must match the block")
        if features.dtype != parameter.dtype:
            raise TypeError("MobileViT features dtype must match the block")
        if not bool(torch.isfinite(features).all().item()):
            raise ValueError("MobileViT features must be finite")
        return features

    def unfold(self, features: Tensor) -> Tensor:
        features = self._validate_features(features, width=self.transformer_width)
        batch, width, height, spatial_width = features.shape
        patch_rows = height // self.patch_height
        patch_columns = spatial_width // self.patch_width
        return (
            features.reshape(
                batch,
                width,
                patch_rows,
                self.patch_height,
                patch_columns,
                self.patch_width,
            )
            .permute(0, 3, 5, 2, 4, 1)
            .reshape(
                batch * self.patch_height * self.patch_width,
                patch_rows * patch_columns,
                width,
            )
        )

    def fold(
        self,
        groups: Tensor,
        *,
        batch: int,
        height: int,
        spatial_width: int,
    ) -> Tensor:
        batch = _efficient_positive_integer(batch, "batch")
        height = _efficient_positive_integer(height, "height")
        spatial_width = _efficient_positive_integer(spatial_width, "spatial_width")
        if height % self.patch_height != 0:
            raise ValueError("MobileViT fold height must be divisible by patch_height")
        if spatial_width % self.patch_width != 0:
            raise ValueError("MobileViT fold width must be divisible by patch_width")
        if not isinstance(groups, Tensor):
            raise TypeError("MobileViT groups must be a tensor")
        if groups.ndim != 3:
            raise ValueError("MobileViT groups must be rank-3")
        if not torch.is_floating_point(groups):
            raise TypeError("MobileViT groups must use a floating dtype")
        if any(dimension == 0 for dimension in groups.shape):
            raise ValueError("MobileViT groups must have non-empty dimensions")
        parameter = self.projection.weight
        if groups.device != parameter.device:
            raise ValueError("MobileViT groups device must match the block")
        if groups.dtype != parameter.dtype:
            raise TypeError("MobileViT groups dtype must match the block")
        if groups.shape[2] != self.transformer_width:
            raise ValueError("MobileViT groups width must match the block")
        patch_rows = height // self.patch_height
        patch_columns = spatial_width // self.patch_width
        expected_shape = (
            batch * self.patch_height * self.patch_width,
            patch_rows * patch_columns,
            self.transformer_width,
        )
        if groups.shape != expected_shape:
            raise ValueError(f"MobileViT groups shape must be {expected_shape}")
        if not bool(torch.isfinite(groups).all().item()):
            raise ValueError("MobileViT groups must be finite")
        return (
            groups.reshape(
                batch,
                self.patch_height,
                self.patch_width,
                patch_rows,
                patch_columns,
                self.transformer_width,
            )
            .permute(0, 5, 3, 1, 4, 2)
            .reshape(batch, self.transformer_width, height, spatial_width)
        )

    def forward(self, features: Tensor) -> tuple[Tensor, Tensor]:
        features = self._validate_features(features)
        batch, _, height, spatial_width = features.shape
        local = self.local(features)
        projected = self.projection(local)
        groups = self.unfold(projected)
        mixed = self.transformer(groups)
        folded = self.fold(
            mixed,
            batch=batch,
            height=height,
            spatial_width=spatial_width,
        )
        projected_back = self.post_fold_projection(folded)
        output = self.fusion(torch.cat((features, projected_back), dim=1))
        assert output.shape == features.shape
        return output, groups


class ToyCascadedGroupAttention(nn.Module):
    def __init__(self, width: int, groups: int, *, ffn_expansion: int = 2) -> None:
        super().__init__()
        self.width = _efficient_positive_integer(width, "width")
        self.groups = _efficient_positive_integer(groups, "groups")
        self.ffn_expansion = _efficient_positive_integer(
            ffn_expansion, "ffn_expansion"
        )
        if self.width % self.groups != 0:
            raise ValueError("width must be divisible by groups")
        self.group_width = self.width // self.groups
        hidden = self.ffn_expansion * self.width
        self.ffn_in = nn.Sequential(
            nn.LayerNorm(self.width),
            nn.Linear(self.width, hidden),
            nn.GELU(),
            nn.Linear(hidden, self.width),
        )
        self.heads = nn.ModuleList(
            [
                nn.MultiheadAttention(self.group_width, 1, batch_first=True)
                for _ in range(self.groups)
            ]
        )
        self.output_projection = nn.Linear(self.width, self.width)
        self.ffn_out = nn.Sequential(
            nn.LayerNorm(self.width),
            nn.Linear(self.width, hidden),
            nn.GELU(),
            nn.Linear(hidden, self.width),
        )

    def _validate_tokens(self, tokens: object) -> Tensor:
        if not isinstance(tokens, Tensor):
            raise TypeError("EfficientViT-CGA tokens must be a tensor")
        if tokens.ndim != 3:
            raise ValueError("EfficientViT-CGA tokens must be rank-3")
        if not torch.is_floating_point(tokens):
            raise TypeError("EfficientViT-CGA tokens must use a floating dtype")
        if any(dimension == 0 for dimension in tokens.shape):
            raise ValueError("EfficientViT-CGA tokens must have non-empty dimensions")
        if tokens.shape[-1] != self.width:
            raise ValueError("EfficientViT-CGA token width must match the block")
        parameter = self.output_projection.weight
        if tokens.device != parameter.device:
            raise ValueError("EfficientViT-CGA tokens device must match the block")
        if tokens.dtype != parameter.dtype:
            raise TypeError("EfficientViT-CGA tokens dtype must match the block")
        if not bool(torch.isfinite(tokens).all().item()):
            raise ValueError("EfficientViT-CGA tokens must be finite")
        return tokens

    def cascade(self, tokens: Tensor) -> tuple[Tensor, ...]:
        tokens = self._validate_tokens(tokens)
        outputs: list[Tensor] = []
        previous: Tensor | None = None
        chunks = tokens.split(self.group_width, dim=-1)
        assert len(chunks) == self.groups
        for chunk, attention in zip(chunks, self.heads):
            current = chunk if previous is None else chunk + previous
            previous, _ = attention(
                current,
                current,
                current,
                need_weights=False,
            )
            outputs.append(previous)
        return tuple(outputs)

    def forward(self, tokens: Tensor) -> Tensor:
        tokens = self._validate_tokens(tokens)
        prepared = tokens + self.ffn_in(tokens)
        grouped = self.cascade(prepared)
        concatenated = torch.cat(grouped, dim=-1)
        assert concatenated.shape == tokens.shape
        projected = prepared + self.output_projection(concatenated)
        output = projected + self.ffn_out(projected)
        assert output.shape == tokens.shape
        return output


def _positive_half_up_retention(patch_count: int, keep_ratio: float) -> int:
    if isinstance(keep_ratio, bool) or not isinstance(keep_ratio, (int, float)):
        raise TypeError("keep_ratio must be a finite number greater than 0 and at most 1")
    if not math.isfinite(keep_ratio) or not 0.0 < keep_ratio <= 1.0:
        raise ValueError("keep_ratio must be a finite number greater than 0 and at most 1")
    return max(1, math.floor(patch_count * keep_ratio + 0.5))


def dynamic_prune_tokens(
    tokens: Tensor,
    scorer: nn.Module,
    keep_ratio: float,
    *,
    training: bool = False,
) -> tuple[Tensor, Tensor]:
    if not isinstance(tokens, Tensor):
        raise TypeError("DynamicViT tokens must be a tensor")
    if tokens.ndim != 3:
        raise ValueError("DynamicViT tokens must be rank-3")
    if not torch.is_floating_point(tokens):
        raise TypeError("DynamicViT tokens must use a floating dtype")
    if tokens.shape[0] == 0 or tokens.shape[2] == 0:
        raise ValueError("DynamicViT tokens must have non-empty dimensions")
    if tokens.shape[1] < 2:
        raise ValueError("DynamicViT tokens must contain CLS and patch tokens")
    if not bool(torch.isfinite(tokens).all().item()):
        raise ValueError("DynamicViT tokens must be finite")
    if not isinstance(scorer, nn.Module):
        raise TypeError("DynamicViT scorer must be an nn.Module")
    if not isinstance(training, bool):
        raise TypeError("DynamicViT training must be a boolean")

    cls_token, patches = tokens[:, :1], tokens[:, 1:]
    keep = _positive_half_up_retention(patches.shape[1], keep_ratio)
    parameters = list(scorer.parameters())
    if parameters:
        parameter = parameters[0]
        if parameter.device != tokens.device:
            raise ValueError("DynamicViT scorer device must match tokens")
        if parameter.dtype != tokens.dtype:
            raise TypeError("DynamicViT scorer dtype must match tokens")
    scores = scorer(patches)
    expected_shape = (*patches.shape[:2], 1)
    if not isinstance(scores, Tensor) or scores.shape != expected_shape:
        raise ValueError("DynamicViT scorer output shape must be [batch, patches, 1]")
    if scores.device != tokens.device:
        raise ValueError("DynamicViT scorer output device must match tokens")
    if scores.dtype != tokens.dtype:
        raise TypeError("DynamicViT scorer output dtype must match tokens")
    if not bool(torch.isfinite(scores).all().item()):
        raise ValueError("DynamicViT scorer output must be finite")
    scores = scores.squeeze(-1)
    top_indices = scores.topk(keep, dim=1).indices

    if training:
        hard_gates = torch.zeros_like(scores).scatter(1, top_indices, 1.0)
        probabilities = torch.sigmoid(scores)
        gates = hard_gates + (probabilities - probabilities.detach())
        gated_patches = patches * gates.unsqueeze(-1)
        return torch.cat((cls_token, gated_patches), dim=1), gates

    indices = top_indices.sort(dim=1).values
    kept = patches.gather(
        1, indices.unsqueeze(-1).expand(-1, -1, patches.shape[-1])
    )
    return torch.cat((cls_token, kept), dim=1), indices


def tome_bipartite_merge(
    tokens: Tensor, sizes: Tensor, merge_count: int
) -> tuple[Tensor, Tensor]:
    if not isinstance(tokens, Tensor):
        raise TypeError("ToMe tokens must be a tensor")
    if tokens.ndim != 3:
        raise ValueError("ToMe tokens must be rank-3")
    if not torch.is_floating_point(tokens):
        raise TypeError("ToMe tokens must use a floating dtype")
    if tokens.shape[0] == 0 or tokens.shape[2] == 0:
        raise ValueError("ToMe tokens must have non-empty dimensions")
    if tokens.shape[1] < 2:
        raise ValueError("ToMe bipartite matching needs at least two tokens")
    if not bool(torch.isfinite(tokens).all().item()):
        raise ValueError("ToMe tokens must be finite")
    if not isinstance(sizes, Tensor):
        raise TypeError("ToMe sizes must be a tensor")
    if sizes.ndim != 2:
        raise ValueError("ToMe sizes must be rank-2")
    if sizes.shape != tokens.shape[:2]:
        raise ValueError("ToMe sizes shape must match [batch, tokens]")
    if not torch.is_floating_point(sizes):
        raise TypeError("ToMe sizes must use a floating dtype")
    if sizes.device != tokens.device:
        raise ValueError("ToMe sizes device must match tokens")
    if sizes.dtype != tokens.dtype:
        raise TypeError("ToMe sizes dtype must match tokens")
    if not bool(torch.isfinite(sizes).all().item()):
        raise ValueError("ToMe sizes must be finite")
    if not bool((sizes > 0).all().item()):
        raise ValueError("ToMe sizes must be positive")
    if isinstance(merge_count, bool) or not isinstance(merge_count, int):
        raise TypeError("merge_count must be a non-negative integer")

    source = tokens[:, 0::2]
    destination = tokens[:, 1::2]
    source_sizes = sizes[:, 0::2]
    destination_sizes = sizes[:, 1::2]
    if merge_count < 0 or merge_count > source.shape[1]:
        raise ValueError("merge_count must not exceed the source set")

    scores = F.normalize(source, dim=-1) @ F.normalize(
        destination, dim=-1
    ).transpose(1, 2)
    assert scores.shape == (
        tokens.shape[0],
        source.shape[1],
        destination.shape[1],
    )
    best_scores, best_destinations = scores.max(dim=-1)
    chosen_sources = best_scores.topk(merge_count, dim=-1).indices
    outputs: list[Tensor] = []
    output_sizes: list[Tensor] = []
    for batch_index in range(tokens.shape[0]):
        chosen = set(chosen_sources[batch_index].tolist())
        kept_indices = torch.tensor(
            [index for index in range(source.shape[1]) if index not in chosen],
            dtype=torch.long,
            device=tokens.device,
        )
        kept_sources = source[batch_index].index_select(0, kept_indices)
        kept_sizes = source_sizes[batch_index].index_select(0, kept_indices)
        merged_destinations: list[Tensor] = []
        merged_destination_sizes: list[Tensor] = []
        for destination_index in range(destination.shape[1]):
            total_size = destination_sizes[batch_index, destination_index]
            weighted = (
                destination[batch_index, destination_index] * total_size
            )
            for source_index in chosen:
                if int(best_destinations[batch_index, source_index]) != destination_index:
                    continue
                source_size = source_sizes[batch_index, source_index]
                weighted = weighted + source[batch_index, source_index] * source_size
                total_size = total_size + source_size
            merged_destinations.append(weighted / total_size)
            merged_destination_sizes.append(total_size)
        outputs.append(
            torch.cat((kept_sources, torch.stack(merged_destinations)), dim=0)
        )
        output_sizes.append(
            torch.cat((kept_sizes, torch.stack(merged_destination_sizes)), dim=0)
        )

    merged = torch.stack(outputs)
    merged_sizes = torch.stack(output_sizes)
    assert merged.shape[:2] == merged_sizes.shape
    assert merged.shape[1] == tokens.shape[1] - merge_count
    return merged, merged_sizes


def tome_proportional_attention(tokens: Tensor, sizes: Tensor) -> tuple[Tensor, Tensor]:
    """Identity-projection attention with ToMe's key-size correction."""
    if not isinstance(tokens, Tensor) or tokens.ndim != 3:
        raise ValueError("ToMe attention tokens must be a rank-3 tensor")
    if not isinstance(sizes, Tensor) or sizes.shape != tokens.shape[:2]:
        raise ValueError("ToMe attention sizes must match [batch, tokens]")
    if sizes.device != tokens.device or sizes.dtype != tokens.dtype:
        raise ValueError("ToMe attention sizes must match token device and dtype")
    if not bool(torch.isfinite(tokens).all().item()):
        raise ValueError("ToMe attention tokens must be finite")
    if not bool(torch.isfinite(sizes).all().item()) or not bool(
        (sizes > 0).all().item()
    ):
        raise ValueError("ToMe attention sizes must be finite and positive")

    logits = tokens @ tokens.transpose(1, 2) / math.sqrt(tokens.shape[-1])
    logits = logits + sizes.log().unsqueeze(1)
    probabilities = logits.softmax(dim=-1)
    return probabilities @ tokens, logits


def run_mobilevit_demo(config: ToyConfig) -> ShapeReport:
    vision = TinyViT(config).cpu()
    block = ToyMobileViTBlock(config.width, config.heads).cpu()
    images = synthetic_images(config).requires_grad_()
    patches = vision.patch_tokens(images)
    feature_map = patches.transpose(1, 2).reshape(
        config.batch_size,
        config.width,
        config.grid_size,
        config.grid_size,
    )
    output, unfolded = block(feature_map)
    assert output.shape == feature_map.shape
    assert unfolded.shape == (
        config.batch_size * block.patch_height * block.patch_width,
        (config.grid_size // block.patch_height)
        * (config.grid_size // block.patch_width),
        config.width,
    )
    loss = output.square().mean()
    assert bool(torch.isfinite(loss).item())
    loss.backward()
    for gradient in (
        images.grad,
        vision.patch_embed.proj.weight.grad,
        block.local[0].weight.grad,
        block.projection.weight.grad,
        block.transformer.self_attn.in_proj_weight.grad,
        block.transformer.linear1.weight.grad,
        block.post_fold_projection.weight.grad,
        block.fusion.weight.grad,
    ):
        assert gradient is not None
        assert bool(torch.isfinite(gradient).all().item())
        assert torch.count_nonzero(gradient) > 0
    return {
        "feature_map": tuple(output.shape),
        "unfolded": tuple(unfolded.shape),
    }


def run_efficientvit_cga_demo(config: ToyConfig) -> ShapeReport:
    model = ToyCascadedGroupAttention(
        config.width,
        config.heads,
        ffn_expansion=2,
    ).cpu()
    tokens = torch.randn(
        config.batch_size,
        config.patch_tokens,
        config.width,
        requires_grad=True,
    )
    output = model(tokens)
    assert output.shape == tokens.shape

    baseline_groups = model.cascade(tokens.detach())
    perturbed = tokens.detach().clone()
    perturbed[..., : model.group_width] += 1
    changed_groups = model.cascade(perturbed)
    for index in range(1, model.groups):
        assert not torch.equal(baseline_groups[index], changed_groups[index])

    loss = output.square().mean()
    assert bool(torch.isfinite(loss).item())
    loss.backward()
    gradients = [
        tokens.grad,
        model.ffn_in[1].weight.grad,
        model.ffn_in[3].weight.grad,
        model.output_projection.weight.grad,
        model.ffn_out[1].weight.grad,
        model.ffn_out[3].weight.grad,
    ]
    for head in model.heads:
        gradients.extend((head.in_proj_weight.grad, head.out_proj.weight.grad))
    for gradient in gradients:
        assert gradient is not None
        assert bool(torch.isfinite(gradient).all().item())
        assert torch.count_nonzero(gradient) > 0
    return {"input": tuple(tokens.shape), "output": tuple(output.shape)}


def run_dynamicvit_demo(config: ToyConfig) -> ShapeReport:
    vision = TinyViT(config).cpu()
    if len(vision.blocks) < 2:
        raise ValueError("DynamicViT demo needs at least two encoder blocks")
    images = synthetic_images(config).requires_grad_()
    with torch.no_grad():
        baseline, _ = vision.forward_features(images.detach())
    assert baseline.shape == (
        config.batch_size,
        config.patch_tokens + 1,
        config.width,
    )

    positioned = vision.add_positions(
        vision.prepend_tokens(vision.patch_tokens(images))
    )
    early = vision.blocks[0](positioned)
    scorer = nn.Linear(config.width, 1).cpu()
    training_tokens, gates = dynamic_prune_tokens(
        early, scorer, 0.5, training=True
    )
    assert training_tokens.shape == early.shape
    assert gates.shape == (config.batch_size, config.patch_tokens)
    later_training = vision.blocks[1](training_tokens)
    for block in vision.blocks[2:]:
        later_training = block(later_training)
    training_loss = later_training.square().mean()
    assert bool(torch.isfinite(training_loss).item())
    training_loss.backward()
    for gradient in (
        scorer.weight.grad,
        vision.patch_embed.proj.weight.grad,
        vision.blocks[0].attn.qkv.weight.grad,
        vision.blocks[1].attn.qkv.weight.grad,
    ):
        assert gradient is not None
        assert bool(torch.isfinite(gradient).all().item())
        assert torch.count_nonzero(gradient) > 0

    later_counts: list[int] = []
    hook = vision.blocks[1].register_forward_pre_hook(
        lambda _module, inputs: later_counts.append(inputs[0].shape[1])
    )
    with torch.no_grad():
        positioned = vision.add_positions(
            vision.prepend_tokens(vision.patch_tokens(images.detach()))
        )
        early = vision.blocks[0](positioned)
        pruned, indices = dynamic_prune_tokens(
            early, scorer, 0.5, training=False
        )
        later_inference = vision.blocks[1](pruned)
    hook.remove()
    kept_patches = _positive_half_up_retention(config.patch_tokens, 0.5)
    assert torch.equal(pruned[:, :1], early[:, :1])
    assert indices.shape == (config.batch_size, kept_patches)
    assert pruned.shape == (
        config.batch_size,
        kept_patches + 1,
        config.width,
    )
    assert later_inference.shape == pruned.shape
    assert later_counts == [kept_patches + 1]
    assert later_counts[0] < early.shape[1]
    return {"indices": tuple(indices.shape), "tokens": tuple(pruned.shape)}


def run_tome_demo(config: ToyConfig) -> ShapeReport:
    tokens = torch.randn(
        config.batch_size,
        8,
        config.width,
        requires_grad=True,
    )
    sizes = torch.ones(
        config.batch_size,
        8,
        requires_grad=True,
    )
    merged, merged_sizes = tome_bipartite_merge(tokens, sizes, merge_count=2)
    assert merged.shape == (config.batch_size, 6, config.width)
    assert merged_sizes.shape == (config.batch_size, 6)
    assert torch.equal(merged_sizes.sum(dim=1), sizes.sum(dim=1))

    _, proportional_logits = tome_proportional_attention(merged, merged_sizes)
    uncorrected_logits = merged @ merged.transpose(1, 2) / math.sqrt(config.width)
    expected_correction = merged_sizes.log().unsqueeze(1).expand_as(
        proportional_logits
    )
    assert torch.allclose(
        proportional_logits - uncorrected_logits, expected_correction
    )

    later = EncoderBlock(config.width, config.heads, config.mlp_ratio).cpu()
    later_counts: list[int] = []
    hook = later.register_forward_pre_hook(
        lambda _module, inputs: later_counts.append(inputs[0].shape[1])
    )
    output = later(merged, merged_sizes)
    hook.remove()
    assert output.shape == merged.shape
    assert later_counts == [6]
    loss = output.square().mean()
    assert bool(torch.isfinite(loss).item())
    loss.backward()
    for gradient in (
        tokens.grad,
        sizes.grad,
        later.attn.qkv.weight.grad,
        later.mlp.fc1.weight.grad,
    ):
        assert gradient is not None
        assert bool(torch.isfinite(gradient).all().item())
        assert torch.count_nonzero(gradient) > 0
    return {"sizes": tuple(merged_sizes.shape), "tokens": tuple(merged.shape)}


DEMOS["mobilevit"] = run_mobilevit_demo
DEMOS["efficientvit-cga"] = run_efficientvit_cga_demo
DEMOS["dynamicvit"] = run_dynamicvit_demo
DEMOS["tome"] = run_tome_demo


def run_demo(name: str, config: ToyConfig) -> ShapeReport:
    if name not in DEMOS:
        raise KeyError(f"unknown demo: {name}")
    return DEMOS[name](config)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description="Mechanism-faithful, synthetic CPU examples for the ViT lineage."
    )
    parser.add_argument(
        "--demo",
        choices=["all", *DEMOS.keys()],
        default="all",
    )
    parser.add_argument("--seed", type=int, default=7)
    args = parser.parse_args(argv)

    names = list(DEMOS) if args.demo == "all" else [args.demo]
    config = ToyConfig()
    for name in names:
        seed_everything(args.seed)
        shapes = run_demo(name, config)
        print(json.dumps({"demo": name, "shapes": shapes}, sort_keys=True))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())