scRep_pretrain/model/foundation_model.py

from __future__ import annotations

from dataclasses import dataclass, fields as dataclass_fields
from contextlib import ExitStack
from typing import Any, Dict, List, Optional, Tuple

import json
import math
import os

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint

try:
    from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
except Exception:  # pragma: no cover - FSDP availability depends on runtime
    FSDP = None


def _dist_is_initialized() -> bool:
    return torch.distributed.is_available() and torch.distributed.is_initialized()


def _leaf_fsdp_modules(module: nn.Module) -> List[nn.Module]:
    if FSDP is None:
        return []

    fsdp_modules: List[nn.Module] = []
    for child in module.modules():
        if not isinstance(child, FSDP):
            continue
        has_fsdp_descendant = any(isinstance(grandchild, FSDP) for grandchild in child.children())
        if not has_fsdp_descendant:
            fsdp_modules.append(child)
    return fsdp_modules


def _json_dump(obj: Any, path: str) -> None:
    os.makedirs(os.path.dirname(path), exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        json.dump(obj, f, ensure_ascii=False, indent=2)


class RmsNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = float(eps)
        self.scale = nn.Parameter(torch.ones(dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        inv = x.pow(2).mean(dim=-1, keepdim=True).add(self.eps).rsqrt()
        return x * inv * self.scale


class SdpaSelfAttention(nn.Module):
    def __init__(self, d_model: int, n_heads: int, dropout: float):
        super().__init__()
        if d_model % n_heads != 0:
            raise ValueError("d_model must be divisible by n_heads")
        self.d_model = int(d_model)
        self.n_heads = int(n_heads)
        self.head_dim = self.d_model // self.n_heads
        self.dropout = float(dropout)

        self.q_proj = nn.Linear(self.d_model, self.d_model, bias=False)
        self.k_proj = nn.Linear(self.d_model, self.d_model, bias=False)
        self.v_proj = nn.Linear(self.d_model, self.d_model, bias=False)
        self.o_proj = nn.Linear(self.d_model, self.d_model, bias=False)

    def forward(self, x: torch.Tensor, *, attn_mask_valid: torch.Tensor) -> torch.Tensor:
        batch_size, seq_len, _ = x.shape
        q = self.q_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)

        if attn_mask_valid.shape != (batch_size, seq_len):
            raise ValueError(
                "attn_mask_valid must have shape [batch_size, seq_len], "
                f"got {tuple(attn_mask_valid.shape)} for x={tuple(x.shape)}"
            )
        # Expression-masked genes remain valid tokens. This mask only excludes
        # right-padding positions from the attention key/value softmax.
        key_mask = attn_mask_valid.to(device=x.device, dtype=torch.bool)[:, None, None, :]
        drop_p = self.dropout if self.training else 0.0
        out = F.scaled_dot_product_attention(
            q,
            k,
            v,
            attn_mask=key_mask,
            dropout_p=drop_p,
            is_causal=False,
        )
        out = out.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
        return self.o_proj(out)


class EncoderBlock(nn.Module):
    def __init__(self, d_model: int, n_heads: int, dropout: float):
        super().__init__()
        self.norm1 = RmsNorm(d_model)
        self.attn = SdpaSelfAttention(d_model, n_heads, dropout)
        self.drop = nn.Dropout(dropout)
        self.norm2 = RmsNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, 4 * d_model, bias=False),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(4 * d_model, d_model, bias=False),
        )

    def forward(self, x: torch.Tensor, *, attn_mask_valid: torch.Tensor) -> torch.Tensor:
        x = x * attn_mask_valid.unsqueeze(-1).to(dtype=x.dtype)
        h = self.norm1(x)
        x = x + self.drop(self.attn(h, attn_mask_valid=attn_mask_valid))
        x = x + self.drop(self.ffn(self.norm2(x)))
        return x * attn_mask_valid.unsqueeze(-1).to(dtype=x.dtype)


@dataclass
class FoundationConfig:
    d_model: int = 768
    n_heads: int = 12
    dropout: float = 0.0
    encoder_layers: int = 12
    gradient_checkpointing: bool = False

    gene_vocab_size: int = 0
    bin_vocab_size: int = 0
    gene_pad_id: int = 0
    bin_pad_id: int = 0
    bin_mask_id: int = 1
    bin_offset: int = 3
    gene_cls_id: int = 1
    expr_embedding_mode: str = "bin"
    expr_rank_mlp_hidden_dim: int = 256

    dino_out_dim: int = 4096
    dino_hidden_dim: int = 2048
    dino_bottleneck_dim: int = 256
    student_temperature: float = 0.10
    teacher_temperature: float = 0.04
    teacher_momentum: float = 0.996
    teacher_momentum_end: float = 1.0
    center_momentum: float = 0.9
    token_center_momentum: float = 0.9
    cross_modal_dino_weight: float = 0.25
    cross_modal_pairing_mode: str = "disabled"
    token_loss_weight: float = 1.0
    token_unmasked_loss_weight: float = 0.0
    inference_use_teacher: bool = True
    teacher_cls_decoder_loss_weight: float = 0.0
    teacher_cls_decoder_hidden_dim: int = 1024
    teacher_cls_decoder_condition_on_obs: bool = False
    teacher_cls_decoder_rank_loss_alpha: float = 0.05
    teacher_cls_decoder_rank_num_pairs: int = 4096
    teacher_student_decoder_distill_weight: float = 0.0
    teacher_student_decoder_loss_mode: str = "distribution"
    teacher_student_decoder_rank_alpha: float = 0.05
    teacher_student_decoder_raw_mmd_weight: float = 0.1
    teacher_student_decoder_raw_rank_weight: float = 0.01
    teacher_student_decoder_raw_mse_weight: float = 1.0


class GeneBinTokenizer(nn.Module):
    def __init__(self, cfg: FoundationConfig):
        super().__init__()
        self.gene_emb = nn.Embedding(cfg.gene_vocab_size, cfg.d_model, padding_idx=cfg.gene_pad_id)
        self.bin_emb = nn.Embedding(cfg.bin_vocab_size, cfg.d_model, padding_idx=cfg.bin_pad_id)
        self.expr_embedding_mode = str(getattr(cfg, "expr_embedding_mode", "bin")).strip().lower()
        if self.expr_embedding_mode not in {"bin", "rank"}:
            raise ValueError(f"Unsupported expr_embedding_mode={self.expr_embedding_mode!r}")
        self.bin_mask_id = int(getattr(cfg, "bin_mask_id", 1))
        self.bin_offset = int(getattr(cfg, "bin_offset", 3))
        rank_hidden_dim = int(getattr(cfg, "expr_rank_mlp_hidden_dim", 256))
        if rank_hidden_dim <= 0:
            raise ValueError("expr_rank_mlp_hidden_dim must be > 0")
        self.rank_mlp: Optional[nn.Sequential]
        self.rank_mask_emb: Optional[nn.Parameter]
        if self.expr_embedding_mode == "rank":
            self.rank_mlp = nn.Sequential(
                nn.Linear(1, rank_hidden_dim, bias=True),
                nn.SiLU(),
                nn.Linear(rank_hidden_dim, cfg.d_model, bias=False),
            )
            self.rank_mask_emb = nn.Parameter(torch.empty(cfg.d_model))
            nn.init.normal_(self.rank_mask_emb, mean=0.0, std=0.02)
        else:
            self.rank_mlp = None
            self.rank_mask_emb = None

    def _rank_expression_emb(self, expr_ids: torch.Tensor) -> torch.Tensor:
        if self.rank_mlp is None:
            raise RuntimeError("rank_mlp is not initialized for non-rank expr_embedding_mode")
        out = self.gene_emb.weight.new_zeros((*expr_ids.shape, self.gene_emb.embedding_dim))
        positive_mask = expr_ids >= int(self.bin_offset)
        if bool(positive_mask.any().item()):
            rank_values = (expr_ids[positive_mask] - int(self.bin_offset) + 1).to(dtype=out.dtype).unsqueeze(-1)
            rank_emb = self.rank_mlp(rank_values)
            if rank_emb.dtype != out.dtype:
                rank_emb = rank_emb.to(dtype=out.dtype)
            out[positive_mask] = rank_emb
        mask_mask = expr_ids == int(self.bin_mask_id)
        if bool(mask_mask.any().item()):
            if self.rank_mask_emb is None:
                raise RuntimeError("rank_mask_emb is not initialized for rank expr_embedding_mode")
            out[mask_mask] = self.rank_mask_emb.to(dtype=out.dtype)
        return out

    def forward(
        self,
        gene_ids: torch.Tensor,
        bin_ids: torch.Tensor,
        attn_mask_valid: torch.Tensor,
        *,
        use_bin_embeddings: bool,
    ) -> torch.Tensor:
        x = self.gene_emb(gene_ids)
        if bool(use_bin_embeddings):
            if self.expr_embedding_mode == "rank":
                x = x + self._rank_expression_emb(bin_ids)
            else:
                x = x + self.bin_emb(bin_ids)
        return x * attn_mask_valid.unsqueeze(-1).to(dtype=x.dtype)


class EncoderBackbone(nn.Module):
    def __init__(self, cfg: FoundationConfig):
        super().__init__()
        self.gradient_checkpointing = bool(getattr(cfg, "gradient_checkpointing", False))
        self.tokenizer = GeneBinTokenizer(cfg)
        self.encoder = nn.ModuleList(
            [EncoderBlock(cfg.d_model, cfg.n_heads, cfg.dropout) for _ in range(cfg.encoder_layers)]
        )
        self.encoder_norm = RmsNorm(cfg.d_model)

    def forward(
        self,
        gene_ids: torch.Tensor,
        bin_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        *,
        use_bin_embeddings: bool,
    ) -> torch.Tensor:
        x = self.tokenizer(
            gene_ids,
            bin_ids,
            attention_mask,
            use_bin_embeddings=use_bin_embeddings,
        )
        for blk in self.encoder:
            if self.gradient_checkpointing and self.training and torch.is_grad_enabled():
                def _block_forward(hidden: torch.Tensor, *, block: EncoderBlock = blk) -> torch.Tensor:
                    return block(hidden, attn_mask_valid=attention_mask)

                x = checkpoint(_block_forward, x, use_reentrant=False)
            else:
                x = blk(x, attn_mask_valid=attention_mask)
        x = self.encoder_norm(x)
        return x * attention_mask.unsqueeze(-1).to(dtype=x.dtype)


class DinoHead(nn.Module):
    def __init__(self, cfg: FoundationConfig):
        super().__init__()
        hidden_dim = int(cfg.dino_hidden_dim)
        bottleneck_dim = int(cfg.dino_bottleneck_dim)
        out_dim = int(cfg.dino_out_dim)
        if hidden_dim <= 0 or bottleneck_dim <= 0 or out_dim <= 0:
            raise ValueError("DINO head dimensions must all be > 0")
        self.mlp = nn.Sequential(
            nn.Linear(cfg.d_model, hidden_dim, bias=False),
            nn.GELU(),
            nn.Linear(hidden_dim, hidden_dim, bias=False),
            nn.GELU(),
            nn.Linear(hidden_dim, bottleneck_dim, bias=False),
        )
        self.norm = nn.LayerNorm(bottleneck_dim)
        self.prototypes = nn.Linear(bottleneck_dim, out_dim, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.mlp(x)
        x = self.norm(x)
        x = F.normalize(x, dim=-1)
        return self.prototypes(x)


class TeacherCLSExpressionDecoder(nn.Module):
    def __init__(self, cfg: FoundationConfig):
        super().__init__()
        hidden_dim = int(getattr(cfg, "teacher_cls_decoder_hidden_dim", 1024))
        if hidden_dim <= 0:
            raise ValueError("teacher_cls_decoder_hidden_dim must be > 0")
        self.condition_on_obs = bool(getattr(cfg, "teacher_cls_decoder_condition_on_obs", False))
        self.cls_proj = nn.Sequential(
            nn.Linear(cfg.d_model, hidden_dim, bias=False),
            nn.LayerNorm(hidden_dim),
            nn.SiLU(),
        )
        self.obs_proj: Optional[nn.Sequential]
        if self.condition_on_obs:
            self.obs_proj = nn.Sequential(
                nn.Linear(2, hidden_dim, bias=False),
                nn.LayerNorm(hidden_dim),
                nn.SiLU(),
            )
        else:
            self.obs_proj = None
        self.gene_proj = nn.Sequential(
            nn.Linear(cfg.d_model, hidden_dim, bias=False),
            nn.LayerNorm(hidden_dim),
            nn.SiLU(),
        )
        self.out_proj = nn.Sequential(
            nn.Linear(2 * hidden_dim, hidden_dim, bias=False),
            nn.SiLU(),
            nn.Linear(hidden_dim, 1, bias=True),
        )

    def forward(
        self,
        cls_state: torch.Tensor,
        gene_query_states: torch.Tensor,
        obs_features: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        cls_hidden = self.cls_proj(cls_state)
        if self.condition_on_obs:
            if obs_features is None:
                raise ValueError("teacher_cls_decoder_condition_on_obs=True requires obs_features")
            cls_hidden = cls_hidden + self.obs_proj(
                obs_features.to(device=cls_hidden.device, dtype=cls_hidden.dtype)
            )
        cls_hidden = cls_hidden.unsqueeze(1).expand(-1, gene_query_states.shape[1], -1)
        gene_hidden = self.gene_proj(gene_query_states)
        return self.out_proj(torch.cat([cls_hidden, gene_hidden], dim=-1)).squeeze(-1)


class CellContrastivePredictModel(nn.Module):
    def __init__(self, cfg: FoundationConfig):
        super().__init__()
        self.cfg = cfg
        if cfg.encoder_layers < 1:
            raise ValueError("encoder_layers must be >= 1")
        if cfg.n_heads < 1:
            raise ValueError("n_heads must be >= 1")
        if cfg.d_model <= 0:
            raise ValueError("d_model must be > 0")
        if cfg.gene_vocab_size <= cfg.gene_cls_id:
            raise ValueError("gene_vocab_size too small for gene_cls_id")
        if cfg.bin_vocab_size <= 0:
            raise ValueError("bin_vocab_size must be > 0")
        if not (0.0 < float(cfg.student_temperature)):
            raise ValueError("student_temperature must be > 0")
        if not (0.0 < float(cfg.teacher_temperature)):
            raise ValueError("teacher_temperature must be > 0")
        if not (0.0 <= float(cfg.token_loss_weight)):
            raise ValueError("token_loss_weight must be >= 0")
        if float(getattr(cfg, "token_unmasked_loss_weight", 0.0)) < 0.0:
            raise ValueError("token_unmasked_loss_weight must be >= 0")
        if not (0.0 <= float(getattr(cfg, "teacher_cls_decoder_loss_weight", 0.0))):
            raise ValueError("teacher_cls_decoder_loss_weight must be >= 0")
        if int(getattr(cfg, "teacher_cls_decoder_hidden_dim", 1024)) <= 0:
            raise ValueError("teacher_cls_decoder_hidden_dim must be > 0")
        if not isinstance(getattr(cfg, "teacher_cls_decoder_condition_on_obs", False), bool):
            raise ValueError("teacher_cls_decoder_condition_on_obs must be a bool")
        if float(getattr(cfg, "teacher_cls_decoder_rank_loss_alpha", 0.05)) < 0.0:
            raise ValueError("teacher_cls_decoder_rank_loss_alpha must be >= 0")
        if int(getattr(cfg, "teacher_cls_decoder_rank_num_pairs", 4096)) < 0:
            raise ValueError("teacher_cls_decoder_rank_num_pairs must be >= 0")
        if float(getattr(cfg, "teacher_student_decoder_distill_weight", 0.0)) < 0.0:
            raise ValueError("teacher_student_decoder_distill_weight must be >= 0")
        loss_mode = str(getattr(cfg, "teacher_student_decoder_loss_mode", "distribution")).strip().lower()
        if loss_mode not in {"distribution", "mse"}:
            raise ValueError("teacher_student_decoder_loss_mode must be one of {'distribution', 'mse'}")
        cfg.teacher_student_decoder_loss_mode = loss_mode
        if float(getattr(cfg, "teacher_student_decoder_rank_alpha", 0.05)) < 0.0:
            raise ValueError("teacher_student_decoder_rank_alpha must be >= 0")
        if float(getattr(cfg, "teacher_student_decoder_raw_mmd_weight", 0.1)) < 0.0:
            raise ValueError("teacher_student_decoder_raw_mmd_weight must be >= 0")
        if float(getattr(cfg, "teacher_student_decoder_raw_rank_weight", 0.01)) < 0.0:
            raise ValueError("teacher_student_decoder_raw_rank_weight must be >= 0")
        if float(getattr(cfg, "teacher_student_decoder_raw_mse_weight", 1.0)) < 0.0:
            raise ValueError("teacher_student_decoder_raw_mse_weight must be >= 0")
        if str(getattr(cfg, "expr_embedding_mode", "bin")).strip().lower() not in {"bin", "rank"}:
            raise ValueError("expr_embedding_mode must be one of {'bin', 'rank'}")
        if int(getattr(cfg, "expr_rank_mlp_hidden_dim", 256)) <= 0:
            raise ValueError("expr_rank_mlp_hidden_dim must be > 0")
        cross_modal_pairing_mode = str(getattr(cfg, "cross_modal_pairing_mode", "disabled")).strip().lower()
        if cross_modal_pairing_mode not in {"disabled", "symmetric", "bin_to_gene_only"}:
            raise ValueError("cross_modal_pairing_mode must be one of {'disabled', 'symmetric', 'bin_to_gene_only'}")
        cfg.cross_modal_pairing_mode = cross_modal_pairing_mode

        self.student_backbone = EncoderBackbone(cfg)
        self.teacher_backbone = EncoderBackbone(cfg)
        self.student_head = DinoHead(cfg)
        self.teacher_head = DinoHead(cfg)
        self.student_ibot_head = DinoHead(cfg)
        self.teacher_ibot_head = DinoHead(cfg)
        self.student_cls_expression_decoder: Optional[TeacherCLSExpressionDecoder] = None
        self.teacher_cls_expression_decoder: Optional[TeacherCLSExpressionDecoder] = None
        self._configure_teacher_cls_expression_decoder(
            loss_weight=float(getattr(cfg, "teacher_cls_decoder_loss_weight", 0.0)),
            hidden_dim=int(getattr(cfg, "teacher_cls_decoder_hidden_dim", 1024)),
        )

        self.register_buffer("center", torch.zeros(1, int(cfg.dino_out_dim)))
        self.register_buffer("token_center", torch.zeros(1, int(cfg.dino_out_dim)))
        self.register_buffer("last_teacher_momentum", torch.tensor(float(cfg.teacher_momentum), dtype=torch.float32))

        self._sync_teacher_from_student()
        self._freeze_teacher()

    def _configure_teacher_cls_expression_decoder(
        self,
        *,
        loss_weight: float,
        hidden_dim: int,
        condition_on_obs: Optional[bool] = None,
    ) -> None:
        current_hidden = int(getattr(self.cfg, "teacher_cls_decoder_hidden_dim", hidden_dim))
        current_condition = bool(getattr(self.cfg, "teacher_cls_decoder_condition_on_obs", False))
        next_condition = current_condition if condition_on_obs is None else bool(condition_on_obs)
        distill_enabled = float(getattr(self.cfg, "teacher_student_decoder_distill_weight", 0.0)) > 0.0
        self.cfg.teacher_cls_decoder_loss_weight = float(loss_weight)
        self.cfg.teacher_cls_decoder_hidden_dim = int(hidden_dim)
        self.cfg.teacher_cls_decoder_condition_on_obs = bool(next_condition)
        if float(loss_weight) <= 0.0:
            self.student_cls_expression_decoder = None
            self.teacher_cls_expression_decoder = None
            return
        if distill_enabled:
            if (
                self.student_cls_expression_decoder is None
                or int(current_hidden) != int(hidden_dim)
                or current_condition != bool(next_condition)
            ):
                self.student_cls_expression_decoder = TeacherCLSExpressionDecoder(self.cfg)
            if (
                self.teacher_cls_expression_decoder is None
                or int(current_hidden) != int(hidden_dim)
                or current_condition != bool(next_condition)
            ):
                self.teacher_cls_expression_decoder = TeacherCLSExpressionDecoder(self.cfg)
            self.teacher_cls_expression_decoder.load_state_dict(self.student_cls_expression_decoder.state_dict())
            for param in self.teacher_cls_expression_decoder.parameters():
                param.requires_grad = False
            return
        self.student_cls_expression_decoder = None
        if (
            self.teacher_cls_expression_decoder is None
            or int(current_hidden) != int(hidden_dim)
            or current_condition != bool(next_condition)
        ):
            self.teacher_cls_expression_decoder = TeacherCLSExpressionDecoder(self.cfg)

    def configure_teacher_cls_expression_decoder(
        self,
        *,
        loss_weight: float,
        hidden_dim: int,
        condition_on_obs: Optional[bool] = None,
    ) -> None:
        self._configure_teacher_cls_expression_decoder(
            loss_weight=float(loss_weight),
            hidden_dim=int(hidden_dim),
            condition_on_obs=condition_on_obs,
        )

    def _freeze_teacher(self) -> None:
        for param in self.teacher_backbone.parameters():
            param.requires_grad = False
        for param in self.teacher_head.parameters():
            param.requires_grad = False
        for param in self.teacher_ibot_head.parameters():
            param.requires_grad = False
        if (
            float(getattr(self.cfg, "teacher_student_decoder_distill_weight", 0.0)) > 0.0
            and self.teacher_cls_expression_decoder is not None
        ):
            for param in self.teacher_cls_expression_decoder.parameters():
                param.requires_grad = False

    def _sync_teacher_from_student(self) -> None:
        self.teacher_backbone.load_state_dict(self.student_backbone.state_dict())
        self.teacher_head.load_state_dict(self.student_head.state_dict())
        self.teacher_ibot_head.load_state_dict(self.student_ibot_head.state_dict())
        if self.student_cls_expression_decoder is not None and self.teacher_cls_expression_decoder is not None:
            self.teacher_cls_expression_decoder.load_state_dict(self.student_cls_expression_decoder.state_dict())

    def _pool_cell_repr(self, token_states: torch.Tensor) -> torch.Tensor:
        return token_states[:, 0, :]

    def _pool_instance_repr(self, token_states: torch.Tensor) -> torch.Tensor:
        return self._pool_cell_repr(token_states)

    def _select_backbone(self, backbone: Optional[str] = None) -> EncoderBackbone:
        chosen = backbone
        if chosen is None:
            chosen = "teacher" if bool(getattr(self.cfg, "inference_use_teacher", True)) else "student"
        if chosen == "teacher":
            return self.teacher_backbone
        if chosen == "student":
            return self.student_backbone
        raise ValueError(f"Unknown backbone={backbone!r}")

    def _encode_tokens(
        self,
        gene_ids: torch.Tensor,
        bin_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        *,
        use_bin_embeddings: bool = True,
        backbone: Optional[str] = None,
    ) -> torch.Tensor:
        module = self._select_backbone(backbone)
        return module(
            gene_ids,
            bin_ids,
            attention_mask,
            use_bin_embeddings=use_bin_embeddings,
        )

    def _encode(
        self,
        gene_ids: torch.Tensor,
        bin_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        *,
        use_bin_embeddings: bool = True,
        backbone: Optional[str] = None,
    ) -> torch.Tensor:
        return self._pool_cell_repr(
            self._encode_tokens(
                gene_ids,
                bin_ids,
                attention_mask,
                use_bin_embeddings=use_bin_embeddings,
                backbone=backbone,
            )
        )

    def _encode_view_stack(
        self,
        gene_ids: torch.Tensor,
        bin_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        view_use_bin: torch.Tensor,
        *,
        backbone_name: str,
        head: nn.Module,
        capture_token_views: Optional[Tuple[int, ...]] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor, Dict[int, torch.Tensor]]:
        num_views = int(gene_ids.shape[1])
        pooled_views = []
        logits_views = []
        token_view_set = set(capture_token_views or ())
        captured_token_states: Dict[int, torch.Tensor] = {}
        for view_idx in range(num_views):
            use_bin = bool(view_use_bin[view_idx].item())
            view_attention_mask = attention_mask[:, view_idx, :]
            valid_positions = torch.nonzero(view_attention_mask, as_tuple=False)
            if valid_positions.numel() == 0:
                raise ValueError(f"view {view_idx} contains no valid tokens")
            max_valid_len = int(valid_positions[:, -1].max().item()) + 1
            tokens = self._encode_tokens(
                gene_ids[:, view_idx, :max_valid_len],
                bin_ids[:, view_idx, :max_valid_len],
                view_attention_mask[:, :max_valid_len],
                use_bin_embeddings=use_bin,
                backbone=backbone_name,
            )
            pooled = self._pool_cell_repr(tokens)
            pooled_views.append(pooled)
            logits_views.append(head(pooled))
            if view_idx in token_view_set:
                captured_token_states[view_idx] = tokens
        return torch.stack(pooled_views, dim=1), torch.stack(logits_views, dim=1), captured_token_states

    def _distributed_batch_mean(self, x: torch.Tensor) -> torch.Tensor:
        mean = x.mean(dim=0, keepdim=True)
        if _dist_is_initialized():
            mean = mean.clone()
            torch.distributed.all_reduce(mean, op=torch.distributed.ReduceOp.SUM)
            mean /= float(torch.distributed.get_world_size())
        return mean

    def _distributed_optional_rank_mean(
        self,
        x: Optional[torch.Tensor],
        *,
        output_dim: int,
        device: torch.device,
    ) -> Optional[torch.Tensor]:
        if x is None or int(x.numel()) == 0 or int(x.shape[0]) == 0:
            mean = torch.zeros((1, int(output_dim)), device=device, dtype=torch.float32)
            active = torch.zeros([], device=device, dtype=torch.float32)
        else:
            mean = x.reshape(-1, x.shape[-1]).detach().float().mean(dim=0, keepdim=True)
            active = torch.ones([], device=mean.device, dtype=torch.float32)

        if _dist_is_initialized():
            torch.distributed.all_reduce(mean, op=torch.distributed.ReduceOp.SUM)
            torch.distributed.all_reduce(active, op=torch.distributed.ReduceOp.SUM)

        if float(active.item()) <= 0.0:
            return None
        return mean / active.clamp_min(1.0)

    def _update_center(self, teacher_logits: torch.Tensor) -> None:
        batch_center = self._distributed_batch_mean(teacher_logits.reshape(-1, teacher_logits.shape[-1]).detach())
        batch_center = batch_center.to(device=self.center.device, dtype=self.center.dtype)
        momentum = float(getattr(self.cfg, "center_momentum", 0.9))
        self.center.mul_(momentum).add_(batch_center, alpha=1.0 - momentum)

    def _update_token_center(self, teacher_logits: Optional[torch.Tensor]) -> None:
        batch_center = self._distributed_optional_rank_mean(
            teacher_logits,
            output_dim=int(self.token_center.shape[-1]),
            device=self.token_center.device,
        )
        if batch_center is None:
            return
        batch_center = batch_center.to(device=self.token_center.device, dtype=self.token_center.dtype)
        momentum = float(getattr(self.cfg, "token_center_momentum", getattr(self.cfg, "center_momentum", 0.9)))
        self.token_center.mul_(momentum).add_(batch_center, alpha=1.0 - momentum)

    def _teacher_probs(self, logits: torch.Tensor) -> torch.Tensor:
        temperature = float(getattr(self.cfg, "teacher_temperature", 0.04))
        logits_f = logits.float()
        center_f = self.center.to(device=logits_f.device, dtype=logits_f.dtype)
        return F.softmax((logits_f - center_f) / temperature, dim=-1)

    def _teacher_token_probs(self, logits: torch.Tensor) -> torch.Tensor:
        temperature = float(getattr(self.cfg, "teacher_temperature", 0.04))
        logits_f = logits.float()
        center_f = self.token_center.to(device=logits_f.device, dtype=logits_f.dtype)
        return F.softmax((logits_f - center_f) / temperature, dim=-1)

    def _student_log_probs(self, logits: torch.Tensor) -> torch.Tensor:
        temperature = float(getattr(self.cfg, "student_temperature", 0.10))
        return F.log_softmax(logits.float() / temperature, dim=-1)

    def _cross_entropy_target(self, student_logits: torch.Tensor, teacher_probs: torch.Tensor) -> torch.Tensor:
        student_log_probs = self._student_log_probs(student_logits)
        return (-teacher_probs * student_log_probs).sum(dim=-1).mean()

    @staticmethod
    def _mean_pairwise_abs_1d(values: torch.Tensor) -> torch.Tensor:
        n = int(values.numel())
        if n <= 1:
            return values.new_zeros(())
        sorted_values = values.sort().values
        coeff = torch.arange(n, device=values.device, dtype=values.dtype) * 2 - (n - 1)
        return (2.0 * (coeff * sorted_values).sum()) / float(n * n)

    @staticmethod
    def _mean_cross_abs_1d(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        nx = int(x.numel())
        ny = int(y.numel())
        if nx == 0 or ny == 0:
            return x.new_zeros(())
        y_sorted = y.sort().values
        prefix = torch.cat([y_sorted.new_zeros((1,)), y_sorted.cumsum(dim=0)])
        right = torch.searchsorted(y_sorted, x, right=False)
        left_sum = prefix[right]
        total_sum = prefix[-1]
        right_f = right.to(dtype=x.dtype)
        ny_f = float(ny)
        abs_sum = x * right_f - left_sum + (total_sum - left_sum) - x * (ny_f - right_f)
        return abs_sum.mean() / ny_f

    def _gene_mmd_loss(self, pred_expr: torch.Tensor, target_expr: torch.Tensor, target_mask: torch.Tensor) -> torch.Tensor:
        mask = target_mask.to(dtype=torch.bool)
        losses = []
        for batch_idx in range(int(pred_expr.shape[0])):
            valid = mask[batch_idx]
            if not bool(valid.any().item()):
                continue
            pred = pred_expr[batch_idx, valid].float()
            target = target_expr[batch_idx, valid].to(device=pred.device, dtype=pred.dtype)
            if int(pred.numel()) <= 1:
                losses.append(F.smooth_l1_loss(pred, target, reduction="mean"))
                continue
            cross = self._mean_cross_abs_1d(pred, target)
            pred_pair = self._mean_pairwise_abs_1d(pred)
            target_pair = self._mean_pairwise_abs_1d(target)
            losses.append((2.0 * cross - pred_pair - target_pair).clamp_min(0.0))
        if not losses:
            return pred_expr.new_zeros(())
        return torch.stack(losses).mean()

    def _rank_loss(self, pred_expr: torch.Tensor, target_expr: torch.Tensor, target_mask: torch.Tensor) -> torch.Tensor:
        num_pairs = int(getattr(self.cfg, "teacher_cls_decoder_rank_num_pairs", 4096))
        if num_pairs <= 0:
            return pred_expr.new_zeros(())
        mask = target_mask.to(dtype=torch.bool)
        losses = []
        for batch_idx in range(int(pred_expr.shape[0])):
            valid = mask[batch_idx]
            if int(valid.sum().item()) <= 1:
                continue
            pred = pred_expr[batch_idx, valid].float()
            target = target_expr[batch_idx, valid].to(device=pred.device, dtype=pred.dtype)
            n = int(pred.numel())
            device = pred.device
            expressed = torch.nonzero(target > 0, as_tuple=False).flatten()
            zero = torch.nonzero(target <= 0, as_tuple=False).flatten()
            pair_i = []
            pair_j = []
            ez_pairs = num_pairs // 2 if int(expressed.numel()) > 0 and int(zero.numel()) > 0 else 0
            if ez_pairs > 0:
                pair_i.append(expressed[torch.randint(int(expressed.numel()), (ez_pairs,), device=device)])
                pair_j.append(zero[torch.randint(int(zero.numel()), (ez_pairs,), device=device)])
            random_pairs = num_pairs - ez_pairs
            if random_pairs > 0:
                pair_i.append(torch.randint(n, (random_pairs,), device=device))
                pair_j.append(torch.randint(n, (random_pairs,), device=device))
            i = torch.cat(pair_i, dim=0)
            j = torch.cat(pair_j, dim=0)
            true_diff = target[i] - target[j]
            true_sign = torch.sign(true_diff)
            valid_pair = true_sign != 0
            if not bool(valid_pair.any().item()):
                continue
            pred_diff = pred[i] - pred[j]
            losses.append(F.softplus(-true_sign[valid_pair] * pred_diff[valid_pair]).mean())
        if not losses:
            return pred_expr.new_zeros(())
        return torch.stack(losses).mean()

    def _rank_loss_random_pairs(
        self,
        pred_expr: torch.Tensor,
        target_expr: torch.Tensor,
        target_mask: torch.Tensor,
    ) -> torch.Tensor:
        num_pairs = int(getattr(self.cfg, "teacher_cls_decoder_rank_num_pairs", 4096))
        if num_pairs <= 0:
            return pred_expr.new_zeros(())
        mask = target_mask.to(dtype=torch.bool)
        losses = []
        for batch_idx in range(int(pred_expr.shape[0])):
            valid = mask[batch_idx]
            if int(valid.sum().item()) <= 1:
                continue
            pred = pred_expr[batch_idx, valid].float()
            target = target_expr[batch_idx, valid].to(device=pred.device, dtype=pred.dtype)
            n = int(pred.numel())
            device = pred.device
            i = torch.randint(n, (num_pairs,), device=device)
            j = torch.randint(n, (num_pairs,), device=device)
            true_sign = torch.sign(target[i] - target[j])
            valid_pair = true_sign != 0
            if not bool(valid_pair.any().item()):
                continue
            pred_diff = pred[i] - pred[j]
            losses.append(F.softplus(-true_sign[valid_pair] * pred_diff[valid_pair]).mean())
        if not losses:
            return pred_expr.new_zeros(())
        return torch.stack(losses).mean()

    def _masked_mse_loss(
        self,
        pred_expr: torch.Tensor,
        target_expr: torch.Tensor,
        target_mask: torch.Tensor,
    ) -> torch.Tensor:
        mask = target_mask.to(device=pred_expr.device, dtype=torch.bool)
        if not bool(mask.any().item()):
            return pred_expr.new_zeros(())
        target = target_expr.to(device=pred_expr.device, dtype=pred_expr.dtype)
        return F.mse_loss(pred_expr[mask].float(), target[mask].float(), reduction="mean")

    def _teacher_cls_decoder_loss(
        self,
        *,
        teacher_gene_ids: torch.Tensor,
        teacher_bin_ids: torch.Tensor,
        teacher_attention_mask: torch.Tensor,
        teacher_view_use_bin: torch.Tensor,
        target_gene_ids: torch.Tensor,
        target_expr: torch.Tensor,
        target_mask: torch.Tensor,
        obs_features: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        if self.teacher_cls_expression_decoder is None:
            return torch.zeros([], device=teacher_gene_ids.device, dtype=torch.float32)
        if target_gene_ids.numel() == 0 or not bool(target_mask.any().item()):
            return torch.zeros([], device=teacher_gene_ids.device, dtype=torch.float32)

        decoder_tokens = self._encode_tokens(
            teacher_gene_ids[:, 0, :],
            teacher_bin_ids[:, 0, :],
            teacher_attention_mask[:, 0, :],
            use_bin_embeddings=bool(teacher_view_use_bin[0].item()),
            backbone="student",
        )
        decoder_cls = self._pool_cell_repr(decoder_tokens)
        gene_queries = self.student_backbone.tokenizer.gene_emb(target_gene_ids)
        pred_expr = self.teacher_cls_expression_decoder(decoder_cls, gene_queries, obs_features=obs_features)

        gene_mmd = self._gene_mmd_loss(pred_expr, target_expr, target_mask)
        rank_alpha = float(getattr(self.cfg, "teacher_cls_decoder_rank_loss_alpha", 0.05))
        if rank_alpha <= 0.0:
            return gene_mmd
        return gene_mmd + rank_alpha * self._rank_loss(pred_expr, target_expr, target_mask)

    def _teacher_student_cls_decoder_loss(
        self,
        *,
        teacher_gene_ids: torch.Tensor,
        teacher_bin_ids: torch.Tensor,
        teacher_attention_mask: torch.Tensor,
        teacher_view_use_bin: torch.Tensor,
        student_gene_ids: torch.Tensor,
        student_bin_ids: torch.Tensor,
        student_attention_mask: torch.Tensor,
        student_view_use_bin: torch.Tensor,
        target_gene_ids: torch.Tensor,
        target_expr: torch.Tensor,
        target_mask: torch.Tensor,
        obs_features: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        if self.student_cls_expression_decoder is None or self.teacher_cls_expression_decoder is None:
            return torch.zeros([], device=teacher_gene_ids.device, dtype=torch.float32)
        if target_gene_ids.numel() == 0 or not bool(target_mask.any().item()):
            return torch.zeros([], device=teacher_gene_ids.device, dtype=torch.float32)

        student_tokens = self._encode_tokens(
            student_gene_ids[:, 0, :],
            student_bin_ids[:, 0, :],
            student_attention_mask[:, 0, :],
            use_bin_embeddings=bool(student_view_use_bin[0].item()),
            backbone="student",
        )
        student_cls = self._pool_cell_repr(student_tokens)
        student_gene_queries = self.student_backbone.tokenizer.gene_emb(target_gene_ids)
        student_scores = self.student_cls_expression_decoder(
            student_cls,
            student_gene_queries,
            obs_features=obs_features,
        )

        with torch.no_grad():
            teacher_tokens = self._encode_tokens(
                teacher_gene_ids[:, 0, :],
                teacher_bin_ids[:, 0, :],
                teacher_attention_mask[:, 0, :],
                use_bin_embeddings=bool(teacher_view_use_bin[0].item()),
                backbone="teacher",
            )
            teacher_cls = self._pool_cell_repr(teacher_tokens)
            teacher_gene_queries = self.teacher_backbone.tokenizer.gene_emb(target_gene_ids)
            teacher_scores = self.teacher_cls_expression_decoder(
                teacher_cls,
                teacher_gene_queries,
                obs_features=obs_features,
            )

        distill_weight = float(getattr(self.cfg, "teacher_student_decoder_distill_weight", 0.0))
        loss_mode = str(getattr(self.cfg, "teacher_student_decoder_loss_mode", "distribution")).strip().lower()
        if loss_mode == "mse":
            raw_mse_weight = float(getattr(self.cfg, "teacher_student_decoder_raw_mse_weight", 1.0))
            loss = student_scores.new_zeros(())
            if distill_weight > 0.0:
                loss = loss + distill_weight * self._masked_mse_loss(student_scores, teacher_scores, target_mask)
            if raw_mse_weight > 0.0:
                loss = loss + raw_mse_weight * self._masked_mse_loss(student_scores, target_expr, target_mask)
            return loss
        if loss_mode != "distribution":
            raise ValueError(f"Unsupported teacher_student_decoder_loss_mode={loss_mode!r}")

        rank_alpha = float(getattr(self.cfg, "teacher_student_decoder_rank_alpha", 0.05))
        raw_mmd_weight = float(getattr(self.cfg, "teacher_student_decoder_raw_mmd_weight", 0.1))
        raw_rank_weight = float(getattr(self.cfg, "teacher_student_decoder_raw_rank_weight", 0.01))

        loss = student_scores.new_zeros(())
        if distill_weight > 0.0:
            loss = loss + distill_weight * self._gene_mmd_loss(student_scores, teacher_scores, target_mask)
            if rank_alpha > 0.0:
                loss = loss + distill_weight * rank_alpha * self._rank_loss_random_pairs(student_scores, teacher_scores, target_mask)
        if raw_mmd_weight > 0.0:
            loss = loss + raw_mmd_weight * self._gene_mmd_loss(student_scores, target_expr, target_mask)
        if raw_rank_weight > 0.0:
            loss = loss + raw_rank_weight * self._rank_loss(student_scores, target_expr, target_mask)
        return loss

    def _allow_cross_mode_pair(self, *, teacher_is_bin: bool, student_is_bin: bool) -> bool:
        mode = str(getattr(self.cfg, "cross_modal_pairing_mode", "bin_to_gene_only"))
        if mode == "disabled":
            return False
        if mode == "symmetric":
            return teacher_is_bin != student_is_bin
        if mode == "bin_to_gene_only":
            return bool(teacher_is_bin) and (not bool(student_is_bin))
        raise ValueError(f"Unknown cross_modal_pairing_mode={mode!r}")

    def _pairwise_dino_loss(
        self,
        student_logits: torch.Tensor,
        student_view_use_bin: torch.Tensor,
        teacher_logits: torch.Tensor,
        teacher_view_use_bin: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        same_losses = []
        cross_losses = []
        cross_weight = float(getattr(self.cfg, "cross_modal_dino_weight", 0.0))

        num_teacher_views = int(teacher_logits.shape[1])
        num_student_views = int(student_logits.shape[1])

        for teacher_idx in range(num_teacher_views):
            teacher_probs = self._teacher_probs(teacher_logits[:, teacher_idx, :]).detach()
            teacher_is_bin = bool(teacher_view_use_bin[teacher_idx].item())
            for student_idx in range(num_student_views):
                student_is_bin = bool(student_view_use_bin[student_idx].item())
                pair_loss = self._cross_entropy_target(
                    student_logits[:, student_idx, :],
                    teacher_probs,
                )
                if student_is_bin == teacher_is_bin:
                    same_losses.append(pair_loss)
                elif cross_weight > 0.0 and self._allow_cross_mode_pair(
                    teacher_is_bin=teacher_is_bin,
                    student_is_bin=student_is_bin,
                ):
                    cross_losses.append(pair_loss)

        device = student_logits.device
        same_loss = (
            torch.stack(same_losses).mean()
            if same_losses
            else torch.zeros([], device=device, dtype=torch.float32)
        )
        cross_loss = (
            torch.stack(cross_losses).mean()
            if cross_losses
            else torch.zeros([], device=device, dtype=torch.float32)
        )
        total = same_loss + float(cross_weight) * cross_loss
        return total, same_loss, cross_loss

    def _matched_token_states_for_pair(
        self,
        *,
        teacher_tokens: torch.Tensor,
        teacher_gene_ids: torch.Tensor,
        teacher_attention_mask: torch.Tensor,
        student_tokens: torch.Tensor,
        student_gene_ids: torch.Tensor,
        student_attention_mask: torch.Tensor,
        student_ibot_mask: torch.Tensor,
        use_masked_tokens: bool = True,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        teacher_states: List[torch.Tensor] = []
        student_states: List[torch.Tensor] = []

        batch_size = int(teacher_tokens.shape[0])
        hidden_dim = int(teacher_tokens.shape[-1])
        for batch_idx in range(batch_size):
            teacher_valid = teacher_attention_mask[batch_idx].bool().clone()
            student_attention_valid = student_attention_mask[batch_idx].bool()
            student_mask = student_ibot_mask[batch_idx].bool()
            if use_masked_tokens:
                student_valid = (student_attention_valid & student_mask).clone()
            else:
                student_valid = (student_attention_valid & (~student_mask)).clone()

            if teacher_valid.numel() <= 1 or student_valid.numel() <= 1:
                continue

            teacher_valid[0] = False
            student_valid[0] = False
            if not teacher_valid.any().item() or not student_valid.any().item():
                continue

            teacher_positions = torch.nonzero(teacher_valid, as_tuple=False).squeeze(-1)
            student_positions = torch.nonzero(student_valid, as_tuple=False).squeeze(-1)
            teacher_genes = teacher_gene_ids[batch_idx, teacher_positions]
            student_genes = student_gene_ids[batch_idx, student_positions]

            matches = student_genes.unsqueeze(1) == teacher_genes.unsqueeze(0)
            matched_student = matches.any(dim=1)
            if not matched_student.any().item():
                continue

            matched_student_positions = student_positions[matched_student]
            matched_teacher_rel_pos = matches[matched_student].to(dtype=torch.int64).argmax(dim=1)
            matched_teacher_positions = teacher_positions[matched_teacher_rel_pos]

            teacher_states.append(teacher_tokens[batch_idx, matched_teacher_positions, :])
            student_states.append(student_tokens[batch_idx, matched_student_positions, :])

        if teacher_states:
            return torch.cat(teacher_states, dim=0), torch.cat(student_states, dim=0)

        empty = teacher_tokens.new_empty((0, hidden_dim))
        return empty, empty

    def _pairwise_token_loss(
        self,
        *,
        student_token_views: Dict[int, torch.Tensor],
        student_gene_ids: torch.Tensor,
        student_attention_mask: torch.Tensor,
        student_ibot_mask: torch.Tensor,
        student_view_use_bin: torch.Tensor,
        student_view_is_global: torch.Tensor,
        teacher_token_views: Dict[int, torch.Tensor],
        teacher_gene_ids: torch.Tensor,
        teacher_attention_mask: torch.Tensor,
        teacher_view_use_bin: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
        same_losses = []
        cross_losses = []
        unmasked_same_losses = []
        unmasked_cross_losses = []
        teacher_logits_for_center: List[torch.Tensor] = []
        cross_weight = float(getattr(self.cfg, "cross_modal_dino_weight", 0.0))
        unmasked_weight = float(getattr(self.cfg, "token_unmasked_loss_weight", 0.0))

        for teacher_idx, teacher_tokens in teacher_token_views.items():
            teacher_is_bin = bool(teacher_view_use_bin[teacher_idx].item())
            for student_idx, student_tokens in student_token_views.items():
                if not bool(student_view_is_global[student_idx].item()):
                    continue
                if not student_ibot_mask[:, student_idx, :].any().item():
                    continue

                student_is_bin = bool(student_view_use_bin[student_idx].item())
                if student_is_bin == teacher_is_bin:
                    allow_pair = True
                elif cross_weight > 0.0:
                    allow_pair = self._allow_cross_mode_pair(
                        teacher_is_bin=teacher_is_bin,
                        student_is_bin=student_is_bin,
                    )
                else:
                    allow_pair = False
                if not allow_pair:
                    continue

                teacher_states, student_states = self._matched_token_states_for_pair(
                    teacher_tokens=teacher_tokens,
                    teacher_gene_ids=teacher_gene_ids[:, teacher_idx, :],
                    teacher_attention_mask=teacher_attention_mask[:, teacher_idx, :],
                    student_tokens=student_tokens,
                    student_gene_ids=student_gene_ids[:, student_idx, :],
                    student_attention_mask=student_attention_mask[:, student_idx, :],
                    student_ibot_mask=student_ibot_mask[:, student_idx, :],
                    use_masked_tokens=True,
                )
                if teacher_states.shape[0] > 0:
                    teacher_logits = self.teacher_ibot_head(teacher_states)
                    teacher_probs = self._teacher_token_probs(teacher_logits).detach()
                    student_logits = self.student_ibot_head(student_states)
                    pair_loss = self._cross_entropy_target(student_logits, teacher_probs)
                    teacher_logits_for_center.append(teacher_logits.detach())

                    if student_is_bin == teacher_is_bin:
                        same_losses.append(pair_loss)
                    else:
                        cross_losses.append(pair_loss)

                if unmasked_weight > 0.0:
                    unmasked_teacher_states, unmasked_student_states = self._matched_token_states_for_pair(
                        teacher_tokens=teacher_tokens,
                        teacher_gene_ids=teacher_gene_ids[:, teacher_idx, :],
                        teacher_attention_mask=teacher_attention_mask[:, teacher_idx, :],
                        student_tokens=student_tokens,
                        student_gene_ids=student_gene_ids[:, student_idx, :],
                        student_attention_mask=student_attention_mask[:, student_idx, :],
                        student_ibot_mask=student_ibot_mask[:, student_idx, :],
                        use_masked_tokens=False,
                    )
                    if unmasked_teacher_states.shape[0] > 0:
                        unmasked_teacher_logits = self.teacher_ibot_head(unmasked_teacher_states)
                        unmasked_teacher_probs = self._teacher_token_probs(unmasked_teacher_logits).detach()
                        unmasked_student_logits = self.student_ibot_head(unmasked_student_states)
                        unmasked_pair_loss = self._cross_entropy_target(
                            unmasked_student_logits,
                            unmasked_teacher_probs,
                        )
                        if student_is_bin == teacher_is_bin:
                            unmasked_same_losses.append(unmasked_pair_loss)
                        else:
                            unmasked_cross_losses.append(unmasked_pair_loss)

        device = student_gene_ids.device
        same_loss = (
            torch.stack(same_losses).mean()
            if same_losses
            else torch.zeros([], device=device, dtype=torch.float32)
        )
        cross_loss = (
            torch.stack(cross_losses).mean()
            if cross_losses
            else torch.zeros([], device=device, dtype=torch.float32)
        )
        unmasked_same_loss = (
            torch.stack(unmasked_same_losses).mean()
            if unmasked_same_losses
            else torch.zeros([], device=device, dtype=torch.float32)
        )
        unmasked_cross_loss = (
            torch.stack(unmasked_cross_losses).mean()
            if unmasked_cross_losses
            else torch.zeros([], device=device, dtype=torch.float32)
        )
        masked_total = same_loss + float(cross_weight) * cross_loss
        unmasked_total = unmasked_same_loss + float(cross_weight) * unmasked_cross_loss
        total = masked_total + unmasked_weight * unmasked_total
        teacher_logits_center = torch.cat(teacher_logits_for_center, dim=0) if teacher_logits_for_center else None
        self._last_token_unmasked_loss = unmasked_total.detach()
        self._last_token_unmasked_same_mode_loss = unmasked_same_loss.detach()
        self._last_token_unmasked_cross_mode_loss = unmasked_cross_loss.detach()
        return total, same_loss, cross_loss, teacher_logits_center

    def _teacher_momentum_value(self, global_step: int, max_steps: int) -> float:
        start = float(getattr(self.cfg, "teacher_momentum", 0.996))
        end = float(getattr(self.cfg, "teacher_momentum_end", 1.0))
        if max_steps <= 1:
            return start
        progress = min(max(float(global_step), 0.0), float(max_steps - 1)) / float(max_steps - 1)
        cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
        return end - (end - start) * cosine

    def _ema_update_module(self, student_module: nn.Module, teacher_module: nn.Module, momentum: float) -> None:
        with ExitStack() as stack:
            if FSDP is not None:
                for wrapped in _leaf_fsdp_modules(student_module):
                    stack.enter_context(FSDP.summon_full_params(wrapped, recurse=False, writeback=False))
                for wrapped in _leaf_fsdp_modules(teacher_module):
                    stack.enter_context(FSDP.summon_full_params(wrapped, recurse=False, writeback=True))

            student_params = dict(student_module.named_parameters())
            teacher_params = dict(teacher_module.named_parameters())
            if student_params.keys() != teacher_params.keys():
                missing_student = sorted(set(teacher_params.keys()) - set(student_params.keys()))
                missing_teacher = sorted(set(student_params.keys()) - set(teacher_params.keys()))
                raise RuntimeError(
                    "Student/teacher parameter names do not match for EMA update: "
                    f"missing_student={missing_student[:5]}, missing_teacher={missing_teacher[:5]}"
                )

            for name, teacher_param in teacher_params.items():
                student_param = student_params[name]
                if teacher_param.data.shape != student_param.data.shape:
                    raise RuntimeError(
                        f"EMA shape mismatch for {name}: teacher={tuple(teacher_param.shape)} "
                        f"student={tuple(student_param.shape)}"
                    )
                teacher_param.data.mul_(momentum).add_(student_param.data, alpha=1.0 - momentum)

    @torch.no_grad()
    def update_teacher(self, *, global_step: int, max_steps: int) -> float:
        momentum = self._teacher_momentum_value(global_step=global_step, max_steps=max_steps)
        self.last_teacher_momentum.fill_(float(momentum))

        self._ema_update_module(self.student_backbone, self.teacher_backbone, momentum)
        self._ema_update_module(self.student_head, self.teacher_head, momentum)
        self._ema_update_module(self.student_ibot_head, self.teacher_ibot_head, momentum)
        if self.student_cls_expression_decoder is not None and self.teacher_cls_expression_decoder is not None:
            self._ema_update_module(
                self.student_cls_expression_decoder,
                self.teacher_cls_expression_decoder,
                momentum,
            )
        return momentum

    def forward(self, batch: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
        student_gene_ids = batch["student_gene_ids"]
        student_bin_ids = batch["student_bin_ids"]
        student_attention_mask = batch["student_attention_mask"]
        student_ibot_mask = batch["student_ibot_mask"]
        student_view_use_bin = batch["student_view_use_bin"]
        student_view_is_global = batch["student_view_is_global"]

        teacher_gene_ids = batch["teacher_gene_ids"]
        teacher_bin_ids = batch["teacher_bin_ids"]
        teacher_attention_mask = batch["teacher_attention_mask"]
        teacher_view_use_bin = batch["teacher_view_use_bin"]

        student_capture_views = tuple(
            idx for idx in range(int(student_view_is_global.numel())) if bool(student_view_is_global[idx].item())
        )
        teacher_capture_views = tuple(range(int(teacher_view_use_bin.numel())))

        _, student_logits, student_token_views = self._encode_view_stack(
            student_gene_ids,
            student_bin_ids,
            student_attention_mask,
            student_view_use_bin,
            backbone_name="student",
            head=self.student_head,
            capture_token_views=student_capture_views,
        )

        with torch.no_grad():
            _, teacher_logits, teacher_token_views = self._encode_view_stack(
                teacher_gene_ids,
                teacher_bin_ids,
                teacher_attention_mask,
                teacher_view_use_bin,
                backbone_name="teacher",
                head=self.teacher_head,
                capture_token_views=teacher_capture_views,
            )

        loss_dino, loss_same_mode, loss_cross_mode = self._pairwise_dino_loss(
            student_logits,
            student_view_use_bin,
            teacher_logits,
            teacher_view_use_bin,
        )

        token_loss_weight = float(getattr(self.cfg, "token_loss_weight", 1.0))
        if token_loss_weight > 0.0:
            loss_token, loss_token_same_mode, loss_token_cross_mode, teacher_token_center_logits = self._pairwise_token_loss(
                student_token_views=student_token_views,
                student_gene_ids=student_gene_ids,
                student_attention_mask=student_attention_mask,
                student_ibot_mask=student_ibot_mask,
                student_view_use_bin=student_view_use_bin,
                student_view_is_global=student_view_is_global,
                teacher_token_views=teacher_token_views,
                teacher_gene_ids=teacher_gene_ids,
                teacher_attention_mask=teacher_attention_mask,
                teacher_view_use_bin=teacher_view_use_bin,
            )
        else:
            zero = torch.zeros([], device=student_logits.device, dtype=torch.float32)
            loss_token = zero
            loss_token_same_mode = zero
            loss_token_cross_mode = zero
            self._last_token_unmasked_loss = zero.detach()
            self._last_token_unmasked_same_mode_loss = zero.detach()
            self._last_token_unmasked_cross_mode_loss = zero.detach()
            teacher_token_center_logits = None

        decoder_loss_weight = float(getattr(self.cfg, "teacher_cls_decoder_loss_weight", 0.0))
        if decoder_loss_weight > 0.0:
            if float(getattr(self.cfg, "teacher_student_decoder_distill_weight", 0.0)) > 0.0:
                loss_decoder = self._teacher_student_cls_decoder_loss(
                    teacher_gene_ids=teacher_gene_ids,
                    teacher_bin_ids=teacher_bin_ids,
                    teacher_attention_mask=teacher_attention_mask,
                    teacher_view_use_bin=teacher_view_use_bin,
                    student_gene_ids=student_gene_ids,
                    student_bin_ids=student_bin_ids,
                    student_attention_mask=student_attention_mask,
                    student_view_use_bin=student_view_use_bin,
                    target_gene_ids=batch["teacher_cls_decoder_target_gene_ids"],
                    target_expr=batch["teacher_cls_decoder_target_expr"],
                    target_mask=batch["teacher_cls_decoder_target_mask"],
                    obs_features=batch.get("teacher_cls_decoder_obs_features"),
                )
            else:
                loss_decoder = self._teacher_cls_decoder_loss(
                    teacher_gene_ids=teacher_gene_ids,
                    teacher_bin_ids=teacher_bin_ids,
                    teacher_attention_mask=teacher_attention_mask,
                    teacher_view_use_bin=teacher_view_use_bin,
                    target_gene_ids=batch["teacher_cls_decoder_target_gene_ids"],
                    target_expr=batch["teacher_cls_decoder_target_expr"],
                    target_mask=batch["teacher_cls_decoder_target_mask"],
                    obs_features=batch.get("teacher_cls_decoder_obs_features"),
                )
        else:
            loss_decoder = torch.zeros([], device=student_logits.device, dtype=torch.float32)

        total_loss = loss_dino + token_loss_weight * loss_token + decoder_loss_weight * loss_decoder

        with torch.no_grad():
            self._update_center(teacher_logits)
            self._update_token_center(teacher_token_center_logits)

        return {
            "loss": total_loss,
            "loss_dino": loss_dino.detach(),
            "loss_dino_same_mode": loss_same_mode.detach(),
            "loss_dino_cross_mode": loss_cross_mode.detach(),
            "loss_token": loss_token.detach(),
            "loss_token_same_mode": loss_token_same_mode.detach(),
            "loss_token_cross_mode": loss_token_cross_mode.detach(),
            "loss_token_unmasked": getattr(self, "_last_token_unmasked_loss", loss_token.new_zeros(())).detach(),
            "loss_token_unmasked_same_mode": getattr(
                self, "_last_token_unmasked_same_mode_loss", loss_token.new_zeros(())
            ).detach(),
            "loss_token_unmasked_cross_mode": getattr(
                self, "_last_token_unmasked_cross_mode_loss", loss_token.new_zeros(())
            ).detach(),
            "loss_decoder": loss_decoder.detach(),
            "loss_total": total_loss.detach(),
            "loss_contrastive": total_loss.detach(),
            "teacher_momentum": self.last_teacher_momentum.detach(),
        }

    def save_pretrained(
        self,
        save_directory: str,
        *,
        save_safetensors: bool = True,
        state_dict: Optional[Dict[str, torch.Tensor]] = None,
        gene_vocab: Optional[Dict[str, Any]] = None,
        meta_vocab: Optional[Dict[str, Any]] = None,
        manifest: Optional[Any] = None,
        extra_files: Optional[Dict[str, Any]] = None,
    ) -> None:
        os.makedirs(save_directory, exist_ok=True)

        cfg_path = os.path.join(save_directory, "model_config.json")
        _json_dump(self.cfg.__dict__, cfg_path)

        if state_dict is None:
            state_dict = self.state_dict()
        cpu_state = {k: v.detach().cpu() for k, v in state_dict.items()}

        if save_safetensors:
            import safetensors.torch as st

            st.save_file(cpu_state, os.path.join(save_directory, "model.safetensors"))
        else:
            torch.save(cpu_state, os.path.join(save_directory, "pytorch_model.bin"))

        if gene_vocab is not None:
            _json_dump(gene_vocab, os.path.join(save_directory, "gene_vocab.json"))
        if meta_vocab is not None:
            _json_dump(meta_vocab, os.path.join(save_directory, "meta_vocab.json"))
        if manifest is not None:
            _json_dump(manifest, os.path.join(save_directory, "manifest.json"))
        if extra_files:
            for name, obj in extra_files.items():
                _json_dump(obj, os.path.join(save_directory, name))

    @classmethod
    def from_pretrained(
        cls,
        load_directory: str,
        *,
        device: Optional[str] = None,
        dtype: Optional[torch.dtype] = None,
        strict: bool = True,
        load_assets: bool = True,
    ) -> Tuple["CellContrastivePredictModel", Optional[Dict[str, Any]], Optional[Dict[str, Any]], Optional[Any]]:
        cfg_path = os.path.join(load_directory, "model_config.json")
        if not os.path.exists(cfg_path):
            raise FileNotFoundError(f"Missing model_config.json in: {load_directory}")

        with open(cfg_path, "r", encoding="utf-8") as f:
            cfg_dict = json.load(f)

        cfg_fields = {f.name for f in dataclass_fields(FoundationConfig)}
        filtered_cfg = {k: v for k, v in cfg_dict.items() if k in cfg_fields}
        cfg = FoundationConfig(**filtered_cfg)
        model = cls(cfg)

        st_path = os.path.join(load_directory, "model.safetensors")
        pt_path = os.path.join(load_directory, "pytorch_model.bin")
        if os.path.exists(st_path):
            import safetensors.torch as st

            state_dict = st.load_file(st_path, device="cpu")
        elif os.path.exists(pt_path):
            state_dict = torch.load(pt_path, map_location="cpu")
        else:
            raise FileNotFoundError(f"Missing model weights in: {load_directory}")

        if "token_center" not in state_dict:
            state_dict["token_center"] = torch.zeros_like(model.token_center)

        for missing_prefix, source_prefix in (
            ("student_ibot_head.", "student_head."),
            ("teacher_ibot_head.", "teacher_head."),
        ):
            if any(key.startswith(missing_prefix) for key in state_dict):
                continue
            for model_key in model.state_dict().keys():
                if not model_key.startswith(missing_prefix):
                    continue
                source_key = source_prefix + model_key[len(missing_prefix) :]
                if source_key in state_dict:
                    state_dict[model_key] = state_dict[source_key].clone()

        model.load_state_dict(state_dict, strict=strict)
        if dtype is not None:
            model = model.to(dtype=dtype)
        if device is not None:
            model = model.to(device=device)

        gene_vocab = None
        meta_vocab = None
        manifest = None
        if load_assets:
            gene_vocab_path = os.path.join(load_directory, "gene_vocab.json")
            meta_vocab_path = os.path.join(load_directory, "meta_vocab.json")
            manifest_path = os.path.join(load_directory, "manifest.json")
            if os.path.exists(gene_vocab_path):
                with open(gene_vocab_path, "r", encoding="utf-8") as f:
                    gene_vocab = json.load(f)
            if os.path.exists(meta_vocab_path):
                with open(meta_vocab_path, "r", encoding="utf-8") as f:
                    meta_vocab = json.load(f)
            if os.path.exists(manifest_path):
                with open(manifest_path, "r", encoding="utf-8") as f:
                    manifest = json.load(f)
        return model, gene_vocab, meta_vocab, manifest


__all__ = ["CellContrastivePredictModel", "FoundationConfig"]