fastplms/models/esmfold2/modeling_esmfold2.py

"""PyTorch ESMFold2 model: the standard released architecture.

Quickstart::

    from transformers import ESMFold2Model

    model = ESMFold2Model.from_pretrained("biohub/ESMFold2").cuda().eval()
    open("ubq.pdb", "w").write(model.infer_protein_as_pdb("MQIFVKTLTGKT..."))

For multi-chain, ligand, and MSA inputs, use ``model.input_types`` together
with ``model.fold(...)`` or ``model.prepare_structure_input(...)``.
"""

from __future__ import annotations

import gc
import importlib
import importlib.metadata
import math
from contextlib import contextmanager
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any, ClassVar, Literal, cast

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from transformers.modeling_outputs import ModelOutput
from transformers.modeling_utils import PreTrainedModel

from ...attention import get_attn_implementation, set_config_attn_implementation

try:
    from fastplms.models.ttt import FastPLMTestTimeTrainingMixin, TTTConfig
except ModuleNotFoundError as error:
    if error.name != "fastplms":
        raise
    from ..ttt import FastPLMTestTimeTrainingMixin, TTTConfig

from .attention import ESMFold2AttentionMixin
from .configuration_esmfold2 import ESMFold2Config, normalize_esmc_id
from .embedding import ESMFold2EmbeddingMixin
from .esmfold2_constants_esm3 import (
    SEQUENCE_BOS_TOKEN,
    SEQUENCE_EOS_TOKEN,
    SEQUENCE_MASK_TOKEN,
    SEQUENCE_PAD_TOKEN,
    SEQUENCE_STANDARD_AA_MAX_TOKEN,
    SEQUENCE_STANDARD_AA_MIN_TOKEN,
    SEQUENCE_VOCAB,
)
from .modeling_esmfold2_common import (
    CHAR_VOCAB_SIZE,
    MAX_ATOMIC_NUMBER,
    MSA_CONDITIONING_INPUT_NAMES,
    NUM_RES_TYPES,
    DiffusionStructureHead,
    FoldingTrunk,
    InputsEmbedder,
    LanguageModelShim,
    MSAPairWeightedAveraging,
    OuterProductMean,
    ResIdxAsymIdSymIdEntityIdEncoding,
    RowAttentionPooling,
    SwiGLUMLP,
    TriangleMultiplicativeUpdate,
    _categorical_mean,
    _compute_intra_token_idx,
    compute_lm_hidden_states,
    gather_rep_atom_coords,
    gather_token_to_atom,
    maybe_apply_msa_column_masking,
    maybe_subsample_msa,
    validate_kernel_backend,
    validate_msa_conditioning_inputs,
    validate_prepared_auxiliary_inputs,
)

_ESMC_FP8_LINEAR_SUFFIX = ".attn.out_proj"
_ESMC_FP8_EXPECTED_PROJECTIONS = 80
_EPS = 1e-6
_NONPOLYMER_ID = 4

# Default for the triangle, OPM, and pair-transition l^2 operations. Caps peak
# memory so l around 2k folds on an 80 GB GPU (about 76 GB at chunk=128 for
# l=1438;
# chunk=64 leaves headroom for the largest foldbench targets). Override via
# ``model.set_chunk_size(...)``; pass None to disable chunking (faster for
# short l but OOM-prone past approximately 600).
_DEFAULT_CHUNK_SIZE = 64


@dataclass
class ESMFold2Output(ModelOutput):
    """Transformers-compatible output shared by released and experimental folds.

    ``last_hidden_state`` is the final pair representation. When requested,
    ``hidden_states`` contains the token-input representation followed by the
    final pair representation. The structure trunks do not expose normalized
    post-softmax attention tensors, so ``output_attentions=True`` fails
    explicitly instead of returning incomplete data.
    """

    last_hidden_state: Tensor | None = None
    hidden_states: tuple[Tensor, ...] | None = None
    attentions: tuple[Tensor, ...] | None = None
    distogram_logits: Tensor | None = None
    sample_atom_coords: Tensor | None = None
    representative_atom_coords: Tensor | None = None
    atom_pad_mask: Tensor | None = None
    residue_index: Tensor | None = None
    entity_id: Tensor | None = None
    plddt_logits: Tensor | None = None
    plddt: Tensor | None = None
    plddt_per_atom: Tensor | None = None
    plddt_ca: Tensor | None = None
    complex_plddt: Tensor | None = None
    complex_iplddt: Tensor | None = None
    pae_logits: Tensor | None = None
    pae: Tensor | None = None
    pde_logits: Tensor | None = None
    pde: Tensor | None = None
    resolved_logits: Tensor | None = None
    ptm: Tensor | None = None
    iptm: Tensor | None = None
    pair_chains_iptm: Tensor | None = None


def _resolve_structure_output_controls(
    config: ESMFold2Config,
    *,
    output_attentions: bool | None,
    output_hidden_states: bool | None,
    return_dict: bool | None,
) -> tuple[bool, bool]:
    resolved_attentions = (
        config.output_attentions if output_attentions is None else output_attentions
    )
    if resolved_attentions:
        raise NotImplementedError(
            "ESMFold2 does not expose normalized attention tensors from its structure "
            "trunk. output_attentions=True is unsupported."
        )
    resolved_hidden_states = (
        config.output_hidden_states
        if output_hidden_states is None
        else output_hidden_states
    )
    resolved_return_dict = config.use_return_dict if return_dict is None else return_dict
    return bool(resolved_hidden_states), bool(resolved_return_dict)


def _finalize_structure_output(
    output: dict[str, Tensor],
    *,
    token_input_state: Tensor,
    pair_state: Tensor,
    output_hidden_states: bool,
    return_dict: bool,
) -> ESMFold2Output | tuple[Any, ...]:
    model_output = ESMFold2Output(
        last_hidden_state=pair_state,
        hidden_states=(token_input_state, pair_state) if output_hidden_states else None,
        **output,
    )
    return model_output if return_dict else model_output.to_tuple()


class _ESMFold2ESMplusplusAdapter(nn.Module):
    def __init__(self, model: nn.Module) -> None:
        super().__init__()
        self.model = model

    @property
    def config(self):
        return self.model.config

    def set_attn_implementation(self, attn_implementation: str) -> None:
        """Update ESMC through its Transformers-compatible attention API."""

        self.model.set_attn_implementation(attn_implementation)

    def forward(
        self,
        input_ids: Tensor,
        attention_mask: Tensor | None = None,
        sequence_id: Tensor | None = None,
        output_hidden_states: bool | None = None,
        output_attentions: bool | None = None,
        return_dict: bool | None = None,
        compute_sae: bool = True,
        normalize_sae: bool = False,
    ):
        del return_dict, compute_sae, normalize_sae
        output = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            sequence_id=sequence_id,
            output_hidden_states=output_hidden_states,
            output_attentions=output_attentions,
            return_dict=True,
            esmfold2_hidden_states=True,
        )
        if output_hidden_states:
            hidden_states = output.hidden_states
            if hidden_states is None:
                raise RuntimeError("ESM++ did not return requested hidden states.")
            if isinstance(hidden_states, torch.Tensor):
                output.hidden_states = hidden_states
            else:
                output.hidden_states = torch.stack(tuple(hidden_states), dim=0)
        return output


def _load_fastplms_esmplusplus_for_esmfold2(
    esmc_model_path: str,
    attn_backend: str,
    device: torch.device,
    dtype: torch.dtype,
    local_files_only: bool = False,
) -> _ESMFold2ESMplusplusAdapter:
    from fastplms.models.esm_plusplus.modeling_esm_plusplus import (
        ESMplusplusConfig,
        ESMplusplusModel,
    )

    normalized_path = normalize_esmc_id(esmc_model_path)
    source_revision, _ = _manifest_esmc_checkpoint_contract(normalized_path)
    revision_kwargs: dict[str, Any] = {
        "local_files_only": local_files_only,
    }
    if source_revision is not None:
        revision_kwargs["revision"] = source_revision
    esmc_config = ESMplusplusConfig.from_pretrained(normalized_path, **revision_kwargs)
    set_config_attn_implementation(esmc_config, attn_backend)
    load_kwargs: dict[str, Any] = {
        "config": esmc_config,
        "torch_dtype": dtype,
        **revision_kwargs,
    }
    if device.type == "cuda":
        # Device mapping constructs parameters on the destination GPU instead
        # of materializing the 6B backbone in host memory first.
        load_kwargs["device_map"] = {"": str(device)}
    esmc = ESMplusplusModel.from_pretrained(normalized_path, **load_kwargs)
    if device.type != "cuda":
        esmc = esmc.to(device=device, dtype=dtype)
    else:
        loaded_device = next(esmc.parameters()).device
        if loaded_device != device:
            raise RuntimeError(
                f"ESMC loaded on {loaded_device}, expected direct loading on {device}."
            )
    return _ESMFold2ESMplusplusAdapter(esmc).eval()


def _manifest_esmc_checkpoint_contract(
    esmc_model_path: str,
) -> tuple[str | None, dict[str, str]]:
    """Return the immutable manifest identity for a registered ESMC source.

    Local checkpoint directories deliberately return no Hub revision. A known
    Hub repository is always loaded at the revision and file identities in
    ``models.toml`` instead of following a mutable branch.
    """

    normalized_path = normalize_esmc_id(esmc_model_path)
    try:
        if Path(normalized_path).exists():
            return None, {}
    except OSError:
        # A repository ID may be too long or otherwise invalid as a local path.
        pass

    from fastplms.registry import get_model_registry

    registry = get_model_registry()
    backbone_model = registry.families["esmfold2"].backbone_model
    if backbone_model is None:
        raise RuntimeError("families.esmfold2 must declare backbone_model.")
    spec = registry[backbone_model]
    for checkpoint in (spec.fast, spec.official):
        if checkpoint.repo_id == normalized_path:
            return checkpoint.revision, {
                item.path: item.encoded for item in checkpoint.files
            }
    if "/" in normalized_path:
        raise ValueError(
            f"Remote ESMC source {normalized_path!r} is not the manifest-declared "
            f"ESMFold2 backbone {spec.fast.repo_id!r}."
        )
    return None, {}


ESMCPrecision = Literal["auto", "bf16", "fp32", "fp8"]


@dataclass(frozen=True, slots=True)
class ESMCPrecisionStatus:
    """Resolved ESMC precision and the evidence used to choose it."""

    requested: str
    resolved: str
    reason: str
    device: str
    transformer_engine_version: str | None

    def as_dict(self) -> dict[str, str | None]:
        return asdict(self)


def _transformer_engine_version() -> str | None:
    try:
        return importlib.metadata.version("transformer-engine")
    except importlib.metadata.PackageNotFoundError:
        return None


def _load_transformer_engine() -> tuple[Any, Any]:
    """Load Transformer Engine lazily so core imports stay dependency-free."""

    try:
        te = importlib.import_module("transformer_engine.pytorch")
        recipe = importlib.import_module("transformer_engine.common.recipe")
    except (ImportError, OSError, RuntimeError) as error:
        raise RuntimeError(
            f"Transformer Engine could not be imported: {type(error).__name__}: {error}"
        ) from error
    if not hasattr(recipe, "Float8CurrentScaling"):
        raise RuntimeError(
            "Transformer Engine does not expose Float8CurrentScaling, which is "
            "required by the validated ESMC FP8 path."
        )
    return te, recipe


def _te_fp8_capability(device: torch.device) -> tuple[bool, str]:
    """Return whether the validated Transformer Engine FP8 path can run."""

    if device.type != "cuda":
        return False, "FP8 requires direct ESMC loading onto a CUDA device."
    if not torch.cuda.is_available():
        return False, "CUDA is unavailable."
    try:
        major, minor = torch.cuda.get_device_capability(device)
    except (AssertionError, RuntimeError, ValueError) as error:
        return False, f"CUDA capability query failed: {error}"
    if not (major >= 9 or (major == 8 and minor >= 9)):
        return False, f"CUDA capability {major}.{minor} does not support FP8."
    try:
        te, _ = _load_transformer_engine()
    except RuntimeError as error:
        return False, str(error)

    probe = getattr(te, "is_fp8_available", None)
    if probe is None:
        try:
            probe = importlib.import_module("transformer_engine.pytorch.fp8").is_fp8_available
        except (ImportError, AttributeError, OSError, RuntimeError) as error:
            return False, f"Transformer Engine has no usable FP8 probe: {error}"
    try:
        try:
            result = probe(return_reason=True)
        except TypeError:
            result = probe()
    except (OSError, RuntimeError) as error:
        return False, f"Transformer Engine FP8 probe failed: {error}"
    if isinstance(result, tuple):
        available = bool(result[0])
        detail = str(result[1]) if len(result) > 1 and result[1] else ""
    else:
        available = bool(result)
        detail = ""
    if not available:
        return False, detail or "Transformer Engine reports FP8 unavailable."
    return True, (
        "Transformer Engine reports FP8 availability; FastPLMs will convert "
        "the validated ESMC attention output projections."
    )


def _resolve_esmc_precision(requested: str, device: torch.device) -> ESMCPrecisionStatus:
    allowed = {"auto", "bf16", "fp32", "fp8"}
    if requested not in allowed:
        raise ValueError(f"precision must be one of {sorted(allowed)}, got {requested!r}.")
    if requested in {"auto", "bf16", "fp32"}:
        resolved = "bf16" if requested == "auto" else requested
        reason = (
            "Automatic precision defaults to BF16; select esmc_precision='fp8' "
            "explicitly to opt in to the validated Transformer Engine path."
            if requested == "auto"
            else "Precision was selected explicitly."
        )
        return ESMCPrecisionStatus(
            requested=requested,
            resolved=resolved,
            reason=reason,
            device=str(device),
            transformer_engine_version=_transformer_engine_version(),
        )
    available, reason = _te_fp8_capability(device)
    if not available:
        raise RuntimeError(f"esmc_precision='fp8' is unavailable: {reason}")
    return ESMCPrecisionStatus(
        requested=requested,
        resolved="fp8",
        reason=reason,
        device=str(device),
        transformer_engine_version=_transformer_engine_version(),
    )


def _install_esmc_backbone(
    model: Any,
    esmc_model_path: str,
    *,
    precision: str,
    device: str | torch.device | None = None,
    local_files_only: bool = False,
) -> None:
    target_device = torch.device(device) if device is not None else model.device
    if target_device.type == "cuda" and target_device.index is None and torch.cuda.is_available():
        target_device = torch.device("cuda", torch.cuda.current_device())
    model_device = torch.device(model.device)
    if model_device.type == "cuda" and model_device.index is None and torch.cuda.is_available():
        model_device = torch.device("cuda", torch.cuda.current_device())
    if target_device != model_device:
        raise ValueError(
            f"ESMC target device {target_device} must match the ESMFold2 device "
            f"{model_device}. Move ESMFold2 before loading or reloading ESMC."
        )
    status = _resolve_esmc_precision(precision, target_device)
    normalized_source = normalize_esmc_id(esmc_model_path)
    source_revision, source_files = _manifest_esmc_checkpoint_contract(normalized_source)
    attention_implementation = get_attn_implementation(model.config)
    model.config.esmc_attn_backend = attention_implementation
    dtype = torch.float32 if status.resolved == "fp32" else torch.bfloat16
    esmc = _load_fastplms_esmplusplus_for_esmfold2(
        esmc_model_path=esmc_model_path,
        attn_backend=attention_implementation,
        device=target_device,
        dtype=dtype,
        local_files_only=local_files_only,
    )
    if esmc.config.hidden_size != model.config.lm_d_model:
        raise ValueError(
            f"ESMFold2 expected lm_d_model={model.config.lm_d_model}, "
            f"but loaded ESMC hidden_size={esmc.config.hidden_size}."
        )
    if esmc.config.num_hidden_layers != model.config.lm_num_layers:
        raise ValueError(
            f"ESMFold2 expected lm_num_layers={model.config.lm_num_layers}, "
            f"but loaded ESMC num_hidden_layers={esmc.config.num_hidden_layers}."
        )
    esmc.eval().requires_grad_(False)
    fp8_module_paths: tuple[str, ...] = ()
    if status.resolved == "fp8":
        fp8_module_paths = _convert_esmc_attention_outputs_to_te(esmc)
        status = ESMCPrecisionStatus(
            requested=status.requested,
            resolved=status.resolved,
            reason=(
                f"{status.reason} Converted {len(fp8_module_paths)} projections; "
                "canonical checkpoint weights remain BF16."
            ),
            device=status.device,
            transformer_engine_version=status.transformer_engine_version,
        )
    model._esmc_source = normalized_source
    model._esmc_source_revision = source_revision
    model._esmc_source_files = source_files
    model._esmc_local_files_only = local_files_only
    model._esmc_precision_policy = precision
    model._esmc_precision_status = status
    model._esmc_fp8 = status.resolved == "fp8"
    model._esmc_fp8_module_paths = fp8_module_paths
    model.config.esmc_precision = precision
    model._esmc = esmc
    model._ttt_lm_head = None


def _drop_transient_esmc_state(
    module: nn.Module,
    state_dict: dict[str, Tensor],
    prefix: str,
    local_metadata: dict[str, Any],
) -> None:
    """Exclude runtime ESMC/TTT modules from canonical folding checkpoints."""

    del module, local_metadata
    transient_prefixes = (f"{prefix}_esmc.", f"{prefix}_ttt_lm_head.")
    for key in tuple(state_dict):
        if key.startswith(transient_prefixes):
            del state_dict[key]


def _reload_esmc_bf16_for_gradients(model: Any, *, reason: str) -> None:
    """Use BF16 temporarily without overwriting the persisted serving policy."""

    policy = model._esmc_precision_policy
    model.reload_esmc(precision="bf16", device=model.device)
    model._esmc_precision_policy = policy
    model.config.esmc_precision = policy
    status = model._esmc_precision_status
    model._esmc_precision_status = ESMCPrecisionStatus(
        requested=policy,
        resolved="bf16",
        reason=reason,
        device=status.device,
        transformer_engine_version=status.transformer_engine_version,
    )


class PairTransition(nn.Module):
    """LayerNorm + SwiGLU feed-forward residual block on the pair representation."""

    def __init__(self, d_model: int, expansion_ratio: int = 4) -> None:
        super().__init__()
        self.norm = nn.LayerNorm(d_model)
        self.ffn = SwiGLUMLP(d_model, expansion_ratio=expansion_ratio, bias=False)
        self._chunk_size: int | None = _DEFAULT_CHUNK_SIZE

    def set_chunk_size(self, chunk_size: int | None) -> None:
        self._chunk_size = chunk_size

    def forward(self, x: Tensor) -> Tensor:
        if self._chunk_size is None or x.shape[1] <= self._chunk_size:
            return self.ffn(self.norm(x))
        out: list[Tensor] = []
        for s in range(0, x.shape[1], self._chunk_size):
            e = min(s + self._chunk_size, x.shape[1])
            sl = x[:, s:e]
            out.append(self.ffn(self.norm(sl)))
        return torch.cat(out, dim=1)


class ConfidenceHead(nn.Module):
    """Predicts pLDDT, PAE, PDE, resolved-atom probability and distogram bins."""

    boundaries: Tensor

    def __init__(self, config: ESMFold2Config) -> None:
        super().__init__()
        ch = config.confidence_head
        d_single = config.d_single
        d_pair = config.d_pair
        d_inputs = config.inputs.d_inputs

        boundaries = torch.linspace(ch.min_dist, ch.max_dist, ch.distogram_bins - 1)
        self.register_buffer("boundaries", boundaries)
        self.dist_bin_pairwise_embed = nn.Embedding(ch.distogram_bins, d_pair)

        self.s_norm = nn.LayerNorm(d_single)
        self.s_inputs_to_single = nn.Linear(d_inputs, d_single, bias=False)
        self.s_to_z = nn.Linear(d_inputs, d_pair, bias=False)
        self.s_to_z_transpose = nn.Linear(d_inputs, d_pair, bias=False)
        self.s_to_z_prod_in1 = nn.Linear(d_inputs, d_pair, bias=False)
        self.s_to_z_prod_in2 = nn.Linear(d_inputs, d_pair, bias=False)
        self.s_to_z_prod_out = nn.Linear(d_pair, d_pair, bias=False)
        self.s_input_to_s = nn.Linear(d_inputs, d_single, bias=False)
        self.s_inputs_norm = nn.LayerNorm(d_inputs)
        self.z_norm = nn.LayerNorm(d_pair)

        self.row_attention_pooling = RowAttentionPooling(d_pair=d_pair, d_single=d_single)

        pf = ch.folding_trunk
        self.folding_trunk = FoldingTrunk(n_layers=pf.n_layers, d_pair=d_pair, expansion_ratio=4)

        # Heads.
        self.plddt_ln = nn.LayerNorm(d_single)
        max_atoms_per_token = 23
        self.plddt_weight = nn.Parameter(
            torch.zeros(max_atoms_per_token, d_single, ch.num_plddt_bins)
        )

        self.pae_ln = nn.LayerNorm(d_pair)
        self.pae_head = nn.Linear(d_pair, ch.num_pae_bins, bias=False)

        self.pde_ln = nn.LayerNorm(d_pair)
        self.pde_head = nn.Linear(d_pair, ch.num_pde_bins, bias=False)

        self.resolved_ln = nn.LayerNorm(d_single)
        # 2 = resolved logits ([unresolved, resolved]).
        self.resolved_weight = nn.Parameter(torch.zeros(max_atoms_per_token, d_single, 2))

    def set_kernel_backend(self, backend: str | None) -> None:
        self.folding_trunk.set_kernel_backend(backend)

    def set_chunk_size(self, chunk_size: int | None) -> None:
        self.folding_trunk.set_chunk_size(chunk_size)

    @staticmethod
    def _repeat_batch(x: Tensor, num_diffusion_samples: int) -> Tensor:
        return x if num_diffusion_samples == 1 else x.repeat_interleave(num_diffusion_samples, 0)

    @staticmethod
    def _flatten_sample_axis(x: Tensor) -> Tensor:
        if x.ndim == 4:
            b, mult, n, c = x.shape
            return x.reshape(b * mult, n, c)
        return x

    def forward(
        self,
        s_inputs: Tensor,
        z: Tensor,
        x_pred: Tensor,
        distogram_atom_idx: Tensor,
        token_attention_mask: Tensor,
        atom_to_token: Tensor,
        atom_attention_mask: Tensor,
        asym_id: Tensor,
        mol_type: Tensor,
        num_diffusion_samples: int = 1,
        relative_position_encoding: Tensor | None = None,
        token_bonds_encoding: Tensor | None = None,
    ) -> dict[str, Tensor]:
        s_inputs_normed = self.s_inputs_norm(s_inputs)

        z_base = self.z_norm(z)
        if relative_position_encoding is not None:
            z_base = z_base + relative_position_encoding
        if token_bonds_encoding is not None:
            z_base = z_base + token_bonds_encoding
        z_base = z_base + self.s_to_z(s_inputs_normed).unsqueeze(2)
        z_base = z_base + self.s_to_z_transpose(s_inputs_normed).unsqueeze(1)
        z_base = z_base + self.s_to_z_prod_out(
            self.s_to_z_prod_in1(s_inputs_normed)[:, :, None, :]
            * self.s_to_z_prod_in2(s_inputs_normed)[:, None, :, :]
        )

        pair = self._repeat_batch(z_base, num_diffusion_samples)
        x_pred_flat = self._flatten_sample_axis(x_pred)
        atom_to_token_m = self._repeat_batch(atom_to_token, num_diffusion_samples)
        atom_mask_m = self._repeat_batch(atom_attention_mask, num_diffusion_samples)
        rep_idx_m = self._repeat_batch(distogram_atom_idx, num_diffusion_samples).long()
        mask = self._repeat_batch(token_attention_mask, num_diffusion_samples)
        expanded_batch_size = pair.shape[0]

        rep_coords = gather_rep_atom_coords(x_pred_flat, rep_idx_m)
        rep_distances = torch.cdist(
            rep_coords, rep_coords, compute_mode="donot_use_mm_for_euclid_dist"
        )
        distogram_bins = (rep_distances.unsqueeze(-1) > self.boundaries).sum(dim=-1).long()
        pair = pair + self.dist_bin_pairwise_embed(distogram_bins)

        pair_mask = mask[:, :, None].float() * mask[:, None, :].float()

        # FoldingTrunk handles the bf16 cast internally during inference so
        # each block's fused trimul engages. In-place residual avoids an
        # extra fp32 pair allocation.
        with torch.amp.autocast("cuda", enabled=pair.is_cuda, dtype=torch.bfloat16):
            pair_delta = self.folding_trunk(pair, pair_attention_mask=pair_mask)
        pair.add_(pair_delta.float())
        del pair_delta
        single = self.row_attention_pooling(pair, mask)

        atom_mask_f = atom_mask_m.float()
        s_at_atoms = gather_token_to_atom(single, atom_to_token_m)
        s_at_atoms_ln = self.plddt_ln(s_at_atoms)

        intra_idx = _compute_intra_token_idx(atom_to_token_m)
        intra_idx = intra_idx.clamp(max=self.plddt_weight.shape[0] - 1)
        w_plddt = self.plddt_weight[intra_idx]
        plddt_logits = torch.einsum("...c,...cb->...b", s_at_atoms_ln, w_plddt)
        plddt_per_atom = _categorical_mean(plddt_logits, start=0.0, end=1.0)

        sequence_length = single.shape[1]
        plddt_sum = torch.zeros(
            expanded_batch_size,
            sequence_length,
            device=single.device,
            dtype=plddt_per_atom.dtype,
        )
        atom_count = torch.zeros(
            expanded_batch_size,
            sequence_length,
            device=single.device,
            dtype=plddt_per_atom.dtype,
        )
        atom_mask_t = atom_mask_f.to(plddt_per_atom.dtype)
        plddt_sum.scatter_add_(1, atom_to_token_m, plddt_per_atom * atom_mask_t)
        atom_count.scatter_add_(1, atom_to_token_m, atom_mask_t)
        plddt = plddt_sum / atom_count.clamp(min=1e-6)

        complex_plddt = (plddt_per_atom * atom_mask_f).sum(dim=-1) / (
            atom_mask_f.sum(dim=-1) + _EPS
        )

        expanded_type = self._repeat_batch(mol_type, num_diffusion_samples)
        expanded_asym = self._repeat_batch(asym_id, num_diffusion_samples)
        is_ligand = (expanded_type == _NONPOLYMER_ID).float()
        inter_chain = (expanded_asym.unsqueeze(-1) != expanded_asym.unsqueeze(-2)).float()
        near_contact = (rep_distances < 8).float()
        interface_per_token = (near_contact * inter_chain * (1.0 - is_ligand).unsqueeze(-1)).amax(
            dim=-1
        )
        iplddt_weight = torch.where(
            is_ligand.bool(),
            torch.full_like(interface_per_token, 2.0),
            interface_per_token,
        )
        iplddt_weight_atoms = gather_token_to_atom(
            iplddt_weight.unsqueeze(-1), atom_to_token_m
        ).squeeze(-1)
        atom_iplddt_w = atom_mask_f * iplddt_weight_atoms
        complex_iplddt = (plddt_per_atom * atom_iplddt_w).sum(dim=-1) / (
            atom_iplddt_w.sum(dim=-1) + _EPS
        )

        plddt_ca = plddt_per_atom.gather(1, rep_idx_m)

        # PAE
        pae_logits = self.pae_head(self.pae_ln(pair))
        pae = _categorical_mean(pae_logits, start=0.0, end=32.0).detach()

        # PDE
        pde_logits = self.pde_head(self.pde_ln(pair))
        pde = _categorical_mean(pde_logits, start=0.0, end=32.0).detach()

        # Resolved (per-atom binary).
        s_at_atoms_res = self.resolved_ln(s_at_atoms)
        w_res = self.resolved_weight[intra_idx]
        resolved_logits = torch.einsum("...c,...cb->...b", s_at_atoms_res, w_res)

        # pTM / ipTM from pae_logits.
        n_bins = pae_logits.shape[-1]
        bin_width = 32.0 / n_bins
        bin_centers = torch.arange(0.5 * bin_width, 32.0, bin_width, device=pae_logits.device)
        mask_f = mask.float()
        n_residues = mask_f.sum(dim=-1, keepdim=True)
        d0 = 1.24 * (n_residues.clamp(min=19) - 15) ** (1 / 3) - 1.8
        tm_per_bin = 1 / (1 + (bin_centers / d0) ** 2)
        pae_probs = F.softmax(pae_logits, dim=-1)
        tm_expected = (pae_probs * tm_per_bin[:, None, None, :]).sum(dim=-1)

        pair_mask_2d = mask_f.unsqueeze(-1) * mask_f.unsqueeze(-2)
        ptm_per_row = (tm_expected * pair_mask_2d).sum(dim=-1) / (pair_mask_2d.sum(dim=-1) + _EPS)
        ptm = ptm_per_row.max(dim=-1).values

        inter_chain_mask = (
            expanded_asym.unsqueeze(-1) != expanded_asym.unsqueeze(-2)
        ).float() * pair_mask_2d
        iptm_per_row = (tm_expected * inter_chain_mask).sum(dim=-1) / (
            inter_chain_mask.sum(dim=-1) + _EPS
        )
        iptm = iptm_per_row.max(dim=-1).values

        max_chain_id = int(expanded_asym.max().item()) if expanded_batch_size > 0 else 0
        n_chains = max_chain_id + 1
        pair_chains_iptm = torch.zeros(
            expanded_batch_size,
            n_chains,
            n_chains,
            device=tm_expected.device,
            dtype=tm_expected.dtype,
        )
        for c1 in range(n_chains):
            chain_c1 = (expanded_asym == c1).float() * mask_f
            if chain_c1.sum() == 0:
                continue
            for c2 in range(n_chains):
                chain_c2 = (expanded_asym == c2).float() * mask_f
                pair_m = chain_c1.unsqueeze(-1) * chain_c2.unsqueeze(-2)
                denom = pair_m.sum(dim=(-1, -2)) + _EPS
                pair_chains_iptm[:, c1, c2] = (tm_expected * pair_m).sum(dim=(-1, -2)) / denom

        return {
            "plddt_logits": plddt_logits,
            "plddt": plddt.detach(),
            "plddt_per_atom": plddt_per_atom.detach(),
            "plddt_ca": plddt_ca.detach(),
            "complex_plddt": complex_plddt.detach(),
            "complex_iplddt": complex_iplddt.detach(),
            "pae_logits": pae_logits,
            "pae": pae,
            "pde_logits": pde_logits,
            "pde": pde,
            "resolved_logits": resolved_logits,
            "ptm": ptm.detach(),
            "iptm": iptm.detach(),
            "pair_chains_iptm": pair_chains_iptm.detach(),
        }


def _inverse_softplus(value: float) -> float:
    return value + math.log(-math.expm1(-value))


def _convert_esmc_attention_outputs_to_te(module: nn.Module) -> tuple[str, ...]:
    """Replace the 80 ESMC attention output projections with TE linears.

    Converting every ESMC linear compounds FP8 error across the 80-layer
    network. The validated inference path limits FP8 GEMMs to each layer's
    attention output projection. Transformer Engine retains canonical BF16
    parameters and creates runtime quantization workspaces during autocast.
    """

    te, _ = _load_transformer_engine()
    converted: list[str] = []

    def walk(owner: nn.Module, prefix: str = "") -> None:
        for name, child in tuple(owner.named_children()):
            path = f"{prefix}.{name}" if prefix else name
            if isinstance(child, nn.Linear) and path.endswith(_ESMC_FP8_LINEAR_SUFFIX):
                replacement = te.Linear(
                    child.in_features,
                    child.out_features,
                    bias=child.bias is not None,
                    params_dtype=child.weight.dtype,
                    device=child.weight.device,
                )
                with torch.no_grad():
                    replacement.weight.copy_(child.weight)
                    if child.bias is not None:
                        replacement.bias.copy_(child.bias)
                replacement.eval().requires_grad_(False)
                setattr(owner, name, replacement)
                converted.append(path)
            else:
                walk(child, path)

    walk(module)
    if len(converted) != _ESMC_FP8_EXPECTED_PROJECTIONS:
        raise RuntimeError(
            "ESMC FP8 conversion expected exactly "
            f"{_ESMC_FP8_EXPECTED_PROJECTIONS} attention output projections, "
            f"found {len(converted)}."
        )
    return tuple(converted)


@contextmanager
def _lm_precision_context(precision: str, device: torch.device):
    """Apply the resolved ESMC inference precision."""

    if device.type != "cuda" or precision == "fp32":
        yield
        return
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        if precision == "fp8":
            te, recipe = _load_transformer_engine()
            fp8_recipe = recipe.Float8CurrentScaling(
                use_power_2_scales=False,
                fp8_format=recipe.Format.HYBRID,
            )
            with te.autocast(enabled=True, recipe=fp8_recipe):
                yield
        else:
            yield


class ESMFold2Model(
    FastPLMTestTimeTrainingMixin,
    ESMFold2EmbeddingMixin,
    ESMFold2AttentionMixin,
    PreTrainedModel,
):
    """ESMFold2: all-atom structure prediction with an ESMC PLM backbone.

    This is the standard released ESMFold2 architecture (uses a linear-
    recurrent trunk, internally referred to as "parcae").

    Forward kwargs that callers commonly override:

    * ``num_loops`` (default ``config.num_loops``): trunk refinement
      loops.
    * ``num_diffusion_samples`` (default ``config.num_diffusion_samples``):
      parallel structure samples; the confidence head re-runs once per
      sample, so memory scales linearly. Pass ``1`` for cheap inference.
    * ``num_sampling_steps`` (default ``config.structure_head.inference_num_steps``):
      diffusion ODE solver steps. Lower for speed, higher for quality.

    Memory / perf knobs:

    * ``model.set_chunk_size(int|None)``: caps l^2 ops (triangle / OPM /
      pair transition) at this token-axis chunk. Default 64: fits
      l approximately 2k on an 80 GB GPU. Pass ``None`` for faster inference
      when l is below 600.
    * ``model.set_kernel_backend(None | "fused" | "cuequivariance")``:
      select kernel backend (None = reference path).
    """

    config_class = ESMFold2Config
    _keys_to_ignore_on_load_unexpected: ClassVar[list[str]] = [r"\._extra_state$"]

    def __init__(self, config: ESMFold2Config) -> None:
        super().__init__(config)
        d_inputs = config.inputs.d_inputs
        d_pair = config.d_pair

        self.inputs_embedder = InputsEmbedder(config)
        self.z_init_1 = nn.Linear(d_inputs, d_pair, bias=False)
        self.z_init_2 = nn.Linear(d_inputs, d_pair, bias=False)
        self.rel_pos = ResIdxAsymIdSymIdEntityIdEncoding(
            n_relative_residx_bins=config.n_relative_residx_bins,
            n_relative_chain_bins=config.n_relative_chain_bins,
            d_pair=d_pair,
        )
        self.token_bonds = nn.Linear(1, d_pair, bias=False)
        self.language_model = LanguageModelShim(
            d_z=d_pair, d_model=config.lm_d_model, num_layers=config.lm_num_layers
        )
        self._esmc: nn.Module | None = None
        self._esmc_fp8: bool = False
        self._esmc_fp8_module_paths: tuple[str, ...] = ()
        self._esmc_source: str = config.esmc_id
        self._esmc_source_revision: str | None = None
        self._esmc_source_files: dict[str, str] = {}
        self._esmc_local_files_only = False
        self._esmc_precision_policy: str = str(getattr(config, "esmc_precision", "auto"))
        self._esmc_precision_status = ESMCPrecisionStatus(
            requested=self._esmc_precision_policy,
            resolved="unloaded",
            reason="ESMC has not been loaded.",
            device=str(self.device),
            transformer_engine_version=_transformer_engine_version(),
        )
        self._ttt_lm_head: nn.Module | None = None
        self._esmfold2_input_builder: Any | None = None
        self._kernel_backend: str | None = None

        pf = config.folding_trunk
        self.folding_trunk = FoldingTrunk(n_layers=pf.n_layers, d_pair=d_pair, expansion_ratio=4)
        if config.lm_encoder.enabled:
            self.lm_encoder: FoldingTrunk | None = FoldingTrunk(
                n_layers=config.lm_encoder.n_layers, d_pair=d_pair, expansion_ratio=4
            )
        else:
            self.lm_encoder = None

        self.parcae_input_norm = nn.LayerNorm(d_pair)
        self.parcae_log_a = nn.Parameter(torch.zeros(d_pair))
        parcae_decay_init = math.sqrt(1.0 / 5.0)
        parcae_delta_init = -math.log(parcae_decay_init)
        self.parcae_log_delta = nn.Parameter(
            torch.full((d_pair,), _inverse_softplus(parcae_delta_init), dtype=torch.float32)
        )
        self.parcae_b_cont = nn.Parameter(torch.eye(d_pair))
        self.parcae_readout = nn.Linear(d_pair, d_pair, bias=False)
        nn.init.eye_(self.parcae_readout.weight)
        self.parcae_coda = FoldingTrunk(
            n_layers=config.parcae.coda_n_layers, d_pair=d_pair, expansion_ratio=4
        )

        # Heads --------------------------------------------------------------
        self.structure_head = DiffusionStructureHead(config)
        self.distogram_head = nn.Linear(d_pair, config.structure_head.distogram_bins, bias=True)
        self.confidence_head = ConfidenceHead(config)

        msa_cfg = config.msa_encoder
        self.msa_encoder = None
        if msa_cfg.enabled:
            self.msa_encoder = MSAEncoder(
                d_msa=msa_cfg.d_msa,
                d_pair=d_pair,
                d_inputs=d_inputs,
                d_hidden=msa_cfg.d_hidden,
                n_layers=msa_cfg.n_layers,
                n_heads_msa=msa_cfg.n_heads_msa,
                msa_head_width=msa_cfg.msa_head_width,
            )

        self.post_init()
        self._register_state_dict_hook(_drop_transient_esmc_state)
        self.init_ttt({"lora_target_replace_module": "MultiHeadAttention"})

    @property
    def esmc_precision_status(self) -> ESMCPrecisionStatus:
        return self._esmc_precision_status

    def load_esmc(
        self,
        esmc_model_path: str,
        precision: ESMCPrecision = "auto",
        device: str | torch.device | None = None,
        local_files_only: bool = False,
    ) -> None:
        """Load canonical ESMC weights and resolve the inference precision."""

        _install_esmc_backbone(
            self,
            esmc_model_path,
            precision=precision,
            device=device,
            local_files_only=local_files_only,
        )

    def reload_esmc(
        self,
        precision: ESMCPrecision = "auto",
        device: str | torch.device | None = None,
        local_files_only: bool | None = None,
    ) -> None:
        """Reload canonical weights with the requested precision policy."""

        source = self._esmc_source or self.config.esmc_id
        old_esmc = self._esmc
        old_head = self._ttt_lm_head
        self._esmc = None
        self._esmc_fp8 = False
        self._esmc_fp8_module_paths = ()
        self._ttt_lm_head = None
        del old_esmc, old_head
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        self.load_esmc(
            source,
            precision=precision,
            device=device,
            local_files_only=(
                self._esmc_local_files_only
                if local_files_only is None
                else local_files_only
            ),
        )

    def _ensure_ttt_bf16(self) -> None:
        if self._esmc_fp8:
            _reload_esmc_bf16_for_gradients(
                self,
                reason="TTT requires BF16; the persisted serving policy is unchanged.",
            )

    def _ensure_ttt_lm_head(self) -> None:
        self._ensure_ttt_bf16()
        if self._esmc is None:
            raise RuntimeError("ESMFold2 TTT requires load_esmc=True.")
        if self._ttt_lm_head is not None:
            return
        from fastplms.models.esm_plusplus.modeling_esm_plusplus import (
            ESMplusplusConfig,
            ESMplusplusForMaskedLM,
        )

        source = self._esmc_source or self.config.esmc_id
        source_revision = self._esmc_source_revision
        if source_revision is None:
            source_revision, _ = _manifest_esmc_checkpoint_contract(source)
        revision_kwargs: dict[str, Any] = {
            "local_files_only": self._esmc_local_files_only,
        }
        if source_revision is not None:
            revision_kwargs["revision"] = source_revision
        esmc_config = ESMplusplusConfig.from_pretrained(
            source,
            **revision_kwargs,
        )
        set_config_attn_implementation(esmc_config, get_attn_implementation(self.config))
        mlm, loading_info = ESMplusplusForMaskedLM.from_pretrained(
            source,
            config=esmc_config,
            output_loading_info=True,
            **revision_kwargs,
        )
        missing_head_keys = [
            key for key in loading_info["missing_keys"] if key.startswith("sequence_head")
        ]
        if missing_head_keys:
            raise RuntimeError(
                "ESMFold2 TTT could not load a pretrained ESM++ MLM head from "
                f"{source}: missing {missing_head_keys}"
            )
        dtype = next(self._esmc.parameters()).dtype
        mlm = mlm.to(device=self.device, dtype=dtype).eval()
        self._ttt_lm_head = mlm.sequence_head
        self._ttt_lm_head.requires_grad_(False)
        del mlm

    def _ttt_get_trainable_modules(self) -> list[nn.Module]:
        self._ensure_ttt_bf16()
        if self._esmc is None:
            raise RuntimeError("ESMFold2 TTT requires load_esmc=True.")
        return [self._esmc]

    def _ttt_tokenize(
        self,
        seq: str | list[str] | None = None,
        input_ids: torch.Tensor | None = None,
        **kwargs,
    ) -> torch.Tensor:
        del kwargs
        if input_ids is not None:
            return input_ids
        if seq is None:
            raise ValueError("Pass either seq or input_ids for ESMFold2 TTT.")
        sequences = [seq] if isinstance(seq, str) else seq
        if not sequences:
            raise ValueError("ESMFold2 TTT requires at least one protein sequence.")
        token_to_id = {token: idx for idx, token in enumerate(SEQUENCE_VOCAB)}
        encoded = []
        for sequence in sequences:
            token_ids = [SEQUENCE_BOS_TOKEN]
            for amino_acid in sequence:
                token_ids.append(token_to_id[amino_acid if amino_acid in token_to_id else "X"])
            token_ids.append(SEQUENCE_EOS_TOKEN)
            encoded.append(token_ids)
        max_len = max(len(token_ids) for token_ids in encoded)
        input_tensor = torch.full(
            (len(encoded), max_len),
            SEQUENCE_PAD_TOKEN,
            dtype=torch.long,
        )
        for row, token_ids in enumerate(encoded):
            input_tensor[row, : len(token_ids)] = torch.tensor(
                token_ids,
                dtype=torch.long,
            )
        return input_tensor

    def _ttt_mask_token(self) -> int:
        return SEQUENCE_MASK_TOKEN

    def _ttt_padding_token(self) -> int:
        return SEQUENCE_PAD_TOKEN

    def _ttt_replacement_tokens(self, input_ids: torch.Tensor) -> torch.Tensor:
        return torch.arange(
            SEQUENCE_STANDARD_AA_MIN_TOKEN,
            SEQUENCE_STANDARD_AA_MAX_TOKEN,
            device=input_ids.device,
            dtype=input_ids.dtype,
        )

    def _ttt_non_special_mask(self, input_ids: torch.Tensor) -> torch.Tensor:
        return (input_ids >= SEQUENCE_STANDARD_AA_MIN_TOKEN) & (
            input_ids < SEQUENCE_STANDARD_AA_MAX_TOKEN
        )

    def _ttt_predict_logits(
        self,
        batch: torch.Tensor | dict[str, torch.Tensor],
        **kwargs,
    ) -> torch.Tensor:
        del kwargs
        if not isinstance(batch, torch.Tensor):
            raise TypeError("ESMFold2 TTT expects input_ids tensors.")
        self._ensure_ttt_bf16()
        if self._esmc is None:
            raise RuntimeError("ESMFold2 TTT requires load_esmc=True.")
        self._ensure_ttt_lm_head()
        if self._ttt_lm_head is None:
            raise RuntimeError("ESMFold2 TTT MLM head initialization failed.")
        attention_mask = batch.ne(SEQUENCE_PAD_TOKEN)
        output = self._esmc(
            input_ids=batch,
            attention_mask=attention_mask,
            return_dict=True,
            compute_sae=False,
        )
        return self._ttt_lm_head(output.last_hidden_state)

    @classmethod
    def from_pretrained(
        cls, pretrained_model_name_or_path, *args, load_esmc: bool = True, **kwargs
    ):
        if cls is ESMFold2Model and "config" not in kwargs:
            config = ESMFold2Config.from_pretrained(pretrained_model_name_or_path, **kwargs)
            if config.type == "experimental":
                raise ValueError(
                    "FastPLMs ESMFold2 supports the released ESMFold2 and "
                    "ESMFold2-Fast checkpoints. Experimental ESMFold2 configs "
                    "are not part of the self-contained AutoModel package."
                )
            kwargs["config"] = config
        # Pop the precision knob before forwarding to the HF loader.
        esmc_precision = kwargs.pop("esmc_precision", None)
        local_files_only = bool(kwargs.get("local_files_only", False))
        output_loading_info = bool(kwargs.get("output_loading_info", False))
        loaded = super().from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
        if output_loading_info:
            model, loading_info = loaded
        else:
            model = loaded
        if load_esmc:
            model.load_esmc(
                model.config.esmc_id,
                precision=esmc_precision or model.config.esmc_precision,
                local_files_only=local_files_only,
            )
        return (model, loading_info) if output_loading_info else model

    def set_kernel_backend(self, backend: str | None) -> None:
        """Select kernel backend.

        Args:
            backend: ``None`` (reference path), ``"fused"`` (requires the
                unavailable source-built Triton bundle), or
                ``"cuequivariance"`` (requires the ``structure,cueq`` extras
                on a supported Linux CUDA 13 host).
        """
        validate_kernel_backend(backend)
        self.folding_trunk.set_kernel_backend(backend)
        if self.lm_encoder is not None:
            self.lm_encoder.set_kernel_backend(backend)
        self.parcae_coda.set_kernel_backend(backend)
        self.confidence_head.set_kernel_backend(backend)
        self.structure_head.set_kernel_backend(backend)
        self._kernel_backend = backend

    def apply_torch_compile(self, mode: str = "fixed_seqlen", dynamic: bool | None = None) -> None:
        """Compile l^2-heavy blocks.

        ``mode='fixed_seqlen'`` recompiles per l; ``'dynamic_seqlen'`` compiles
        once.

        Does NOT stack with our Triton kernels: call ``set_kernel_backend(None)``
        before compiling.
        """
        if dynamic is None:
            dynamic = mode == "dynamic_seqlen"
        kwargs: dict = {"dynamic": dynamic}

        from .modeling_esmfold2_common import (
            DiffusionModule,
            DiffusionTransformer,
            PairUpdateBlock,
        )

        compile_targets = (
            PairUpdateBlock,
            DiffusionTransformer,
            DiffusionModule,
            MSAEncoderBlock,
        )

        def _maybe_compile(module: nn.Module) -> None:
            if isinstance(module, compile_targets):
                module.forward = torch.compile(module.forward, **kwargs)  # type: ignore[assignment]

        self.apply(_maybe_compile)

    def set_chunk_size(self, chunk_size: int | None) -> None:
        self.folding_trunk.set_chunk_size(chunk_size)
        if self.lm_encoder is not None:
            self.lm_encoder.set_chunk_size(chunk_size)
        self.parcae_coda.set_chunk_size(chunk_size)
        self.confidence_head.set_chunk_size(chunk_size)
        if self.msa_encoder is not None:
            self.msa_encoder.set_chunk_size(chunk_size)

    def _compute_lm_hidden_states(
        self,
        input_ids: Tensor,
        asym_id: Tensor,
        residue_index: Tensor,
        mol_type: Tensor,
        tok_mask: Tensor,
        lm_mask_pct: float = 0.0,
    ) -> Tensor:
        if self._esmc_fp8 and torch.is_grad_enabled():
            _reload_esmc_bf16_for_gradients(
                self,
                reason=(
                    "Gradient-enabled ESMC execution requires BF16; the persisted "
                    "serving policy is unchanged."
                ),
            )
        if self._esmc is None:
            raise RuntimeError("ESMFold2 requires load_esmc=True for LM feature extraction.")
        # Transformer Engine FP8 kernels require l to be a multiple of 16.
        pad_to = 16 if self._esmc_fp8 else None
        with _lm_precision_context(self._esmc_precision_status.resolved, self.device):
            return compute_lm_hidden_states(
                self._esmc,
                input_ids,
                asym_id,
                residue_index,
                mol_type,
                tok_mask,
                pad_to_multiple=pad_to,
                lm_mask_pct=lm_mask_pct,
                mask_token_id=SEQUENCE_MASK_TOKEN,
            )

    def _discretized_dynamics(self) -> tuple[Tensor, Tensor]:
        delta = F.softplus(self.parcae_log_delta)
        a = torch.exp(-delta * torch.exp(self.parcae_log_a))
        b = delta[:, None] * self.parcae_b_cont
        return a, b

    def _init_pair_state(self, ref: Tensor) -> Tensor:
        std = math.sqrt(2.0 / (5.0 * ref.shape[-1]))
        state = torch.empty_like(ref, dtype=torch.float32)
        nn.init.trunc_normal_(state, mean=0.0, std=std, a=-3 * std, b=3 * std)
        return state.to(dtype=ref.dtype)

    def _run_one_loop(
        self,
        z: Tensor,
        z_init: Tensor,
        lm_z: Tensor | None,
        _msa_inputs: dict | None,
        pair_mask: Tensor,
        a: Tensor,
        b_mat: Tensor,
        tok_mask: Tensor,
        total_steps: int,
    ) -> Tensor:
        # Helper method (not inline) so per-iter locals free on return:
        # otherwise leaks about 2 GB of l^2 * c_z data into distogram/sample scope.
        # training=True forces dropout under eval(), matching the per-loop
        # dropout strategy used at train time.
        lm_cfg = self.config.lm_encoder
        _per_loop_lm_dropout = (
            lm_z is not None
            and getattr(lm_cfg, "per_loop_lm_dropout", False)
            and getattr(lm_cfg, "lm_dropout", 0.0) > 0.0
        )
        _lm_dropout_p = getattr(lm_cfg, "lm_dropout", 0.0)

        for _ in range(total_steps):
            if _per_loop_lm_dropout:
                if lm_z is None:
                    raise RuntimeError("Per-loop LM dropout requires LM pair features.")
                lm_z_i: Tensor | None = F.dropout(lm_z, p=_lm_dropout_p, training=True)
            else:
                lm_z_i = lm_z

            refined_lm_z: Tensor | None = None
            if lm_z_i is not None and self.lm_encoder is not None:
                refined_lm_z = self.lm_encoder(
                    lm_z_i.to(z_init.dtype), pair_attention_mask=pair_mask
                )

            z_inject_pair = z_init
            if lm_z_i is not None and self.lm_encoder is None:
                z_inject_pair = z_inject_pair + lm_z_i.to(z_inject_pair.dtype)

            if self.msa_encoder is not None and _msa_inputs is not None:
                msa_i, mask_i, hd_i, dv_i = maybe_subsample_msa(
                    _msa_inputs["msa"],
                    _msa_inputs["msa_attention_mask"],
                    _msa_inputs["has_deletion"],
                    _msa_inputs["deletion_value"],
                    max_depth=_msa_inputs["max_depth"],
                    enabled=_msa_inputs["subsample_enabled"],
                )
                b_msa, m, l_msa = msa_i.shape
                msa_oh = F.one_hot(msa_i.permute(0, 2, 1).long(), num_classes=NUM_RES_TYPES).float()
                msa_attn = (
                    mask_i.permute(0, 2, 1).float()
                    if mask_i is not None
                    else tok_mask[:, :, None].expand(-1, -1, m).float()
                )
                # Bias-free MSAEncoder.embed requires zeroed padding.
                msa_oh = msa_oh * msa_attn.unsqueeze(-1)
                hd = (
                    hd_i.permute(0, 2, 1).float()
                    if hd_i is not None
                    else torch.zeros(b_msa, l_msa, m, device=msa_i.device)
                )
                dv = (
                    dv_i.permute(0, 2, 1).float()
                    if dv_i is not None
                    else torch.zeros(b_msa, l_msa, m, device=msa_i.device)
                )
                msa_pair = self.msa_encoder(
                    x_pair=z_inject_pair,
                    x_inputs=_msa_inputs["x_inputs"],
                    msa_oh=msa_oh,
                    has_deletion=hd,
                    deletion_value=dv,
                    msa_attention_mask=msa_attn,
                ).to(z_inject_pair.dtype)
                z_inject_pair = (
                    msa_pair if self.config.msa_encoder_overwrite else (z_inject_pair + msa_pair)
                )

            if refined_lm_z is not None:
                z_inject_pair = z_inject_pair + refined_lm_z.to(z_inject_pair.dtype)

            injected_pair = self.parcae_input_norm(z_inject_pair)
            z = a * z + F.linear(injected_pair.to(z.dtype), b_mat)
            z = self.folding_trunk(z, pair_attention_mask=pair_mask)

        return z

    def forward(
        self,
        token_index: Tensor,
        residue_index: Tensor,
        asym_id: Tensor,
        sym_id: Tensor,
        entity_id: Tensor,
        mol_type: Tensor,
        res_type: Tensor,
        token_bonds: Tensor,
        token_attention_mask: Tensor,
        ref_pos: Tensor,
        ref_element: Tensor,
        ref_charge: Tensor,
        ref_atom_name_chars: Tensor,
        ref_space_uid: Tensor,
        atom_attention_mask: Tensor,
        atom_to_token: Tensor,
        distogram_atom_idx: Tensor,
        deletion_mean: Tensor | None = None,
        msa: Tensor | None = None,
        has_deletion: Tensor | None = None,
        deletion_value: Tensor | None = None,
        msa_attention_mask: Tensor | None = None,
        input_ids: Tensor | None = None,
        lm_hidden_states: Tensor | None = None,
        num_loops: int | None = None,
        num_diffusion_samples: int | None = None,
        num_sampling_steps: int | None = None,
        lm_mask_pct: float | None = None,
        msa_max_depth: int = 1024,
        msa_column_mask_rate: float = 0.1,
        msa_subsample_at_inference: bool = True,
        early_exit: bool = False,
        noise_scale: float | None = None,
        step_scale: float | None = None,
        max_inference_sigma: float | None = None,
        output_attentions: bool | None = None,
        output_hidden_states: bool | None = None,
        return_dict: bool | None = None,
        pocket_feature: Tensor | None = None,
        gt_coords: Tensor | None = None,
        is_resolved: Tensor | None = None,
        frames_idx: Tensor | None = None,
        disto_cond: Tensor | None = None,
        disto_cond_mask: Tensor | None = None,
    ) -> ESMFold2Output | tuple[Any, ...]:
        output_hidden_states, return_dict = _resolve_structure_output_controls(
            self.config,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        validate_msa_conditioning_inputs(
            self.config,
            msa=msa,
            msa_attention_mask=msa_attention_mask,
            has_deletion=has_deletion,
            deletion_value=deletion_value,
            deletion_mean=deletion_mean,
        )
        validate_prepared_auxiliary_inputs(
            pocket_feature=pocket_feature,
            disto_cond=disto_cond,
            disto_cond_mask=disto_cond_mask,
        )
        del gt_coords, is_resolved, frames_idx
        tok_mask = token_attention_mask
        atm_mask = atom_attention_mask
        disto_idx = distogram_atom_idx

        n_loops: int = num_loops if num_loops is not None else self.config.num_loops
        n_samples: int = (
            num_diffusion_samples
            if num_diffusion_samples is not None
            else self.config.num_diffusion_samples
        )
        total_steps = max(1, n_loops + 1)

        if res_type.dim() == 2:
            res_type_oh = F.one_hot(res_type.long(), num_classes=NUM_RES_TYPES).float()
            res_type_oh = res_type_oh * tok_mask.unsqueeze(-1).float()
        else:
            res_type_oh = res_type.float()

        if msa is not None:
            msa_oh_profile = F.one_hot(msa.long(), num_classes=NUM_RES_TYPES).float()
            if msa_attention_mask is not None:
                mask_f = msa_attention_mask.float().unsqueeze(-1)
                msa_oh_profile = msa_oh_profile * mask_f
                valid_seq_count = msa_attention_mask.float().sum(dim=1).clamp(min=1)
                profile = msa_oh_profile.sum(dim=1) / valid_seq_count.unsqueeze(-1)
            else:
                profile = msa_oh_profile.mean(dim=1)
        else:
            profile = res_type_oh

        if deletion_mean is None:
            deletion_mean = torch.zeros(
                res_type.shape[0], res_type.shape[1], device=res_type.device
            )

        ref_element_oh = F.one_hot(ref_element.long(), num_classes=MAX_ATOMIC_NUMBER).float()
        ref_atom_name_chars_oh = F.one_hot(
            ref_atom_name_chars.long(), num_classes=CHAR_VOCAB_SIZE
        ).float()
        # Bias-free downstream Linears require zeroed padding.
        atm_mask_f = atm_mask.float()
        ref_element_oh = ref_element_oh * atm_mask_f.unsqueeze(-1)
        ref_atom_name_chars_oh = ref_atom_name_chars_oh * atm_mask_f.unsqueeze(-1).unsqueeze(-1)
        atom_to_token = atom_to_token * atm_mask.long()

        use_amp = ref_pos.device.type == "cuda"
        with torch.amp.autocast("cuda", enabled=use_amp, dtype=torch.bfloat16):
            x_inputs = self.inputs_embedder(
                aatype=res_type_oh,
                profile=profile.float(),
                deletion_mean=deletion_mean.float(),
                ref_pos=ref_pos,
                atom_attention_mask=atm_mask,
                ref_space_uid=ref_space_uid,
                ref_charge=ref_charge,
                ref_element=ref_element_oh,
                ref_atom_name_chars=ref_atom_name_chars_oh,
                atom_to_token=atom_to_token,
            )

            z_init = self.z_init_1(x_inputs).unsqueeze(2) + self.z_init_2(x_inputs).unsqueeze(1)

            relative_position_encoding = self.rel_pos(
                residue_index=residue_index,
                asym_id=asym_id,
                sym_id=sym_id,
                entity_id=entity_id,
                token_index=token_index,
            )
            token_bonds_encoding = self.token_bonds(token_bonds.float())
            z_init = z_init + relative_position_encoding + token_bonds_encoding

            if lm_hidden_states is None and input_ids is not None and self._esmc is not None:
                lm_hidden_states = self._compute_lm_hidden_states(
                    input_ids,
                    asym_id,
                    residue_index,
                    mol_type,
                    tok_mask,
                    lm_mask_pct=(self.config.lm_mask_pct if lm_mask_pct is None else lm_mask_pct),
                )
            lm_z: Tensor | None = None
            if lm_hidden_states is not None:
                lm_z = self.language_model(lm_hidden_states.detach())
            del lm_hidden_states

            pair_mask = tok_mask[:, :, None].float() * tok_mask[:, None, :].float()

            z = self._init_pair_state(z_init)

            a, b = self._discretized_dynamics()
            a = a.view(1, 1, 1, -1).to(device=z.device, dtype=z.dtype)
            b_mat = b.to(device=z.device, dtype=z.dtype)

            _msa_inputs: dict | None = None
            if self.msa_encoder is not None and msa is not None:
                msa_attention_mask = maybe_apply_msa_column_masking(
                    msa_attention_mask,
                    msa_column_mask_rate,
                )
                _msa_inputs = dict(
                    x_inputs=x_inputs,
                    msa=msa,
                    msa_attention_mask=msa_attention_mask,
                    has_deletion=has_deletion,
                    deletion_value=deletion_value,
                    max_depth=msa_max_depth,
                    subsample_enabled=msa_subsample_at_inference,
                )

            # Method call (not inline loop) frees per-iteration l^2 * c_z locals.
            z = self._run_one_loop(
                z=z,
                z_init=z_init,
                lm_z=lm_z,
                _msa_inputs=_msa_inputs,
                pair_mask=pair_mask,
                a=a,
                b_mat=b_mat,
                tok_mask=tok_mask,
                total_steps=total_steps,
            )
            del z_init, lm_z, _msa_inputs, a, b_mat

            z = self.parcae_readout(z)
            z = self.parcae_coda(z, pair_attention_mask=pair_mask)

            z = z.float()
        distogram_logits = self.distogram_head(z + z.transpose(-2, -3))

        structure_output = self.structure_head.sample(
            z_trunk=z,
            s_inputs=x_inputs,
            s_trunk=None,
            relative_position_encoding=relative_position_encoding,
            ref_pos=ref_pos,
            ref_charge=ref_charge,
            ref_mask=atm_mask,
            ref_element=ref_element_oh,
            ref_atom_name_chars=ref_atom_name_chars_oh,
            ref_space_uid=ref_space_uid,
            tok_idx=atom_to_token,
            asym_id=asym_id,
            residue_index=residue_index,
            entity_id=entity_id,
            token_index=token_index,
            sym_id=sym_id,
            token_attention_mask=tok_mask,
            num_diffusion_samples=n_samples,
            num_sampling_steps=num_sampling_steps,
            max_inference_sigma=max_inference_sigma,
            noise_scale=noise_scale,
            step_scale=step_scale,
            return_atom_repr=False,
            denoising_early_exit_rmsd=(0.10 if early_exit else None),
        )

        sample_coords = structure_output["sample_atom_coords"]
        if sample_coords is None:
            raise RuntimeError("ESMFold2 structure sampling did not return coordinates.")
        output: dict[str, Tensor] = {"distogram_logits": distogram_logits}
        output["sample_atom_coords"] = sample_coords

        confidence_output = self.confidence_head(
            s_inputs=x_inputs.detach(),
            z=z.detach().float(),
            x_pred=sample_coords.detach(),
            distogram_atom_idx=disto_idx,
            token_attention_mask=tok_mask,
            atom_to_token=atom_to_token,
            atom_attention_mask=atm_mask,
            asym_id=asym_id,
            mol_type=mol_type,
            num_diffusion_samples=n_samples,
            relative_position_encoding=relative_position_encoding.detach(),
            token_bonds_encoding=token_bonds_encoding.detach(),
        )
        output.update(confidence_output)
        output["atom_pad_mask"] = atm_mask.unsqueeze(0) if atm_mask.dim() == 1 else atm_mask
        output["residue_index"] = residue_index
        output["entity_id"] = entity_id
        return _finalize_structure_output(
            output,
            token_input_state=x_inputs,
            pair_state=z,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

    @torch.no_grad()
    def infer_protein(self, seq: str, **forward_kwargs) -> ESMFold2Output:
        from .protein_utils import prepare_protein_features

        if forward_kwargs.pop("return_dict", True) is not True:
            raise ValueError(
                "infer_protein always returns a mapping; return_dict=False is invalid."
            )
        features = prepare_protein_features(seq)
        if not self.config.msa_conditioning:
            for name in MSA_CONDITIONING_INPUT_NAMES:
                features.pop(name, None)
        features = {k: v.to(self.device) for k, v in features.items()}
        return self(**features, **forward_kwargs, return_dict=True)

    @property
    def input_builder(self):
        if self._esmfold2_input_builder is None:
            from .esmfold2_processor import ESMFold2InputBuilder

            self._esmfold2_input_builder = ESMFold2InputBuilder()
        return self._esmfold2_input_builder

    @property
    def input_types(self):
        from . import esmfold2_types

        return esmfold2_types

    def prepare_structure_input(self, input, seed: int | None = None):
        return self.input_builder.prepare_model_input(
            self,
            input,
            seed=seed,
            device=self.device,
        )

    def fold(
        self,
        input,
        *,
        num_loops: int = 3,
        num_sampling_steps: int = 50,
        num_diffusion_samples: int = 1,
        seed: int | None = None,
        noise_scale: float | None = None,
        step_scale: float | None = None,
        max_inference_sigma: int | None = None,
        early_exit: bool = False,
        complex_id: str = "pred",
    ):
        return self.input_builder.fold(
            self,
            input,
            num_loops=num_loops,
            num_sampling_steps=num_sampling_steps,
            num_diffusion_samples=num_diffusion_samples,
            seed=seed,
            noise_scale=noise_scale,
            step_scale=step_scale,
            max_inference_sigma=max_inference_sigma,
            early_exit=early_exit,
            complex_id=complex_id,
        )

    def _fold_protein_no_ttt(
        self,
        sequence: str,
        *,
        chain_id: str = "A",
        msa: Any | None = None,
        msa_path: str | Path | None = None,
        msa_max_sequences: int | None = None,
        num_loops: int = 3,
        num_sampling_steps: int = 50,
        num_diffusion_samples: int = 1,
        seed: int | None = None,
        complex_id: str = "pred",
    ):
        from .esmfold2_types import MSA, ProteinInput, StructurePredictionInput

        if msa is not None and msa_path is not None:
            raise ValueError("Pass at most one of msa or msa_path.")
        if msa_path is not None:
            msa = MSA.from_a3m(msa_path, max_sequences=msa_max_sequences)
        if msa is not None:
            query = str(msa.query).replace("-", "").upper()
            if query != sequence.upper():
                raise ValueError(
                    "MSA query does not match sequence: "
                    f"expected {sequence.upper()!r}, got {query!r}"
                )

        input = StructurePredictionInput(
            sequences=[ProteinInput(id=chain_id, sequence=sequence, msa=msa)]
        )
        return self.fold(
            input,
            num_loops=num_loops,
            num_sampling_steps=num_sampling_steps,
            num_diffusion_samples=num_diffusion_samples,
            seed=seed,
            complex_id=complex_id,
        )

    @staticmethod
    def _ttt_mean_plddt(result) -> float:
        if result.plddt is None:
            raise RuntimeError("ESMFold2 result has no pLDDT tensor.")
        return float(result.plddt.float().mean().item())

    def _ttt_select_result(self, result):
        if isinstance(result, list):
            if not result:
                raise RuntimeError("ESMFold2 fold returned an empty result list.")
            return max(result, key=self._ttt_mean_plddt)
        return result

    def _ttt_eval_step(
        self,
        step: int,
        loss: float,
        seq: str | list[str] | None = None,
        input_ids: torch.Tensor | None = None,
        **kwargs,
    ) -> tuple[dict[str, Any], float | None]:
        del input_ids
        if not isinstance(seq, str):
            raise TypeError("ESMFold2 fold TTT is protein-only and sequence-string only.")
        fold_kwargs = kwargs["fold_kwargs"]
        was_training = self.training
        self.eval()
        try:
            result = self._fold_protein_no_ttt(seq, **fold_kwargs)
        finally:
            self.train(was_training)
        selected = self._ttt_select_result(result)
        plddt = self._ttt_mean_plddt(selected)
        return {
            "step": step,
            "loss": loss,
            "plddt": plddt,
            "result": selected,
        }, plddt

    def fold_protein(
        self,
        sequence: str,
        *,
        chain_id: str = "A",
        msa: Any | None = None,
        msa_path: str | Path | None = None,
        msa_max_sequences: int | None = None,
        num_loops: int = 3,
        num_sampling_steps: int = 50,
        num_diffusion_samples: int = 1,
        seed: int | None = None,
        complex_id: str = "pred",
        ttt: bool = False,
        ttt_config: TTTConfig | dict[str, Any] | None = None,
    ):
        if ttt:
            return self.fold_protein_ttt(
                sequence=sequence,
                chain_id=chain_id,
                msa=msa,
                msa_path=msa_path,
                msa_max_sequences=msa_max_sequences,
                num_loops=num_loops,
                num_sampling_steps=num_sampling_steps,
                num_diffusion_samples=num_diffusion_samples,
                seed=seed,
                complex_id=complex_id,
                ttt_config=ttt_config,
            )
        return self._fold_protein_no_ttt(
            sequence=sequence,
            chain_id=chain_id,
            msa=msa,
            msa_path=msa_path,
            msa_max_sequences=msa_max_sequences,
            num_loops=num_loops,
            num_sampling_steps=num_sampling_steps,
            num_diffusion_samples=num_diffusion_samples,
            seed=seed,
            complex_id=complex_id,
        )

    def fold_protein_ttt(
        self,
        sequence: str,
        *,
        chain_id: str = "A",
        msa: Any | None = None,
        msa_path: str | Path | None = None,
        msa_max_sequences: int | None = None,
        num_loops: int = 3,
        num_sampling_steps: int = 50,
        num_diffusion_samples: int = 1,
        seed: int | None = None,
        complex_id: str = "pred",
        ttt_config: TTTConfig | dict[str, Any] | None = None,
    ):
        self._ensure_ttt_bf16()
        if self._esmc is None:
            raise RuntimeError("ESMFold2 TTT requires load_esmc=True.")
        fold_kwargs = {
            "chain_id": chain_id,
            "msa": msa,
            "msa_path": msa_path,
            "msa_max_sequences": msa_max_sequences,
            "num_loops": num_loops,
            "num_sampling_steps": num_sampling_steps,
            "num_diffusion_samples": num_diffusion_samples,
            "seed": seed,
            "complex_id": complex_id,
        }
        baseline = self._ttt_select_result(self._fold_protein_no_ttt(sequence, **fold_kwargs))
        baseline_plddt = self._ttt_mean_plddt(baseline)
        best_result = baseline
        best_plddt = baseline_plddt
        best_step = 0
        step_plddts = [baseline_plddt]

        cfg = self.ttt_config.merged(ttt_config).merged(
            {"eval_each_step": True, "automatic_best_state_reset": False}
        )
        try:
            metrics = self.ttt(
                seq=sequence,
                ttt_config=cfg,
                fold_kwargs=fold_kwargs,
            )
            for step_metric in metrics["step_metrics"]:
                step_plddt = step_metric["plddt"]
                step_plddts.append(step_plddt)
                if step_plddt > best_plddt:
                    best_plddt = step_plddt
                    best_step = step_metric["step"]
                    best_result = step_metric["result"]
            best_result.ttt_metrics = {
                "losses": metrics["losses"],
                "step_plddts": step_plddts,
                "baseline_plddt": baseline_plddt,
                "best_plddt": best_plddt,
                "best_step": best_step,
            }
            return best_result
        finally:
            if "_ttt_initialized" in self.__dict__ and self._ttt_initialized:
                self.ttt_reset()

    @staticmethod
    def result_to_cif(result) -> str:
        if isinstance(result, list):
            raise TypeError("Pass one MolecularComplexResult at a time.")
        return result.complex.to_mmcif()

    @staticmethod
    def result_to_pdb(result) -> str:
        if isinstance(result, list):
            raise TypeError("Pass one MolecularComplexResult at a time.")
        return result.complex.to_protein_complex().to_pdb_string()

    def save_as_cif(self, result, output_path: str | Path) -> None:
        Path(output_path).write_text(self.result_to_cif(result))

    def save_as_pdb(self, result, output_path: str | Path) -> None:
        Path(output_path).write_text(self.result_to_pdb(result))

    def infer_protein_as_cif(self, seq: str, **forward_kwargs) -> str:
        return self.result_to_cif(self.fold_protein(seq, **forward_kwargs))

    def infer_protein_as_pdb(self, seq: str, **forward_kwargs) -> str:
        return self.result_to_pdb(self.fold_protein(seq, **forward_kwargs))


class MSAEncoderBlock(nn.Module):
    """One MSA encoder block: OPM into pair, MSA pair-weighted averaging, triangle update."""

    def __init__(
        self,
        d_msa: int,
        d_pair: int,
        d_hidden: int,
        n_heads_msa: int,
        msa_head_width: int,
        is_final_block: bool = False,
    ) -> None:
        super().__init__()
        self.is_final_block = is_final_block
        self.outer_product_mean = OuterProductMean(d_msa, d_hidden, d_pair)
        if not is_final_block:
            self.msa_pair_weighted_averaging = MSAPairWeightedAveraging(
                d_msa, d_pair, n_heads_msa, msa_head_width
            )
            self.msa_transition = PairTransition(d_msa, expansion_ratio=4)
        self.tri_mul_out = TriangleMultiplicativeUpdate(dim=d_pair, _outgoing=True)
        self.tri_mul_in = TriangleMultiplicativeUpdate(dim=d_pair, _outgoing=False)
        self.pair_transition = PairTransition(d_pair, expansion_ratio=4)

    def set_chunk_size(self, chunk_size: int | None) -> None:
        self.outer_product_mean.set_chunk_size(chunk_size)
        self.tri_mul_out.set_chunk_size(chunk_size)
        self.tri_mul_in.set_chunk_size(chunk_size)
        if not self.is_final_block:
            self.msa_transition.set_chunk_size(chunk_size)
        self.pair_transition.set_chunk_size(chunk_size)

    def forward(
        self,
        m: Tensor,
        pair: Tensor,
        msa_attention_mask: Tensor,
        pair_attention_mask: Tensor,
    ) -> tuple[Tensor, Tensor]:
        pair = pair + self.outer_product_mean(m, msa_attention_mask)
        if not self.is_final_block:
            m = m + self.msa_pair_weighted_averaging(m, pair, pair_attention_mask)
            m = m + self.msa_transition(m)
        pair = pair + self.tri_mul_out(pair, mask=pair_attention_mask)
        pair = pair + self.tri_mul_in(pair, mask=pair_attention_mask)
        pair = pair + self.pair_transition(pair)
        return m, pair


class MSAEncoder(nn.Module):
    """Stack of [`MSAEncoderBlock`] layers that conditions the pair on an MSA."""

    def __init__(
        self,
        d_msa: int,
        d_pair: int,
        d_inputs: int,
        d_hidden: int = 32,
        n_layers: int = 4,
        n_heads_msa: int = 8,
        msa_head_width: int = 16,
    ) -> None:
        super().__init__()
        self.embed = nn.Linear(35, d_msa, bias=False)
        self.project_inputs = nn.Linear(d_inputs, d_msa, bias=False)
        self.blocks = nn.ModuleList(
            [
                MSAEncoderBlock(
                    d_msa=d_msa,
                    d_pair=d_pair,
                    d_hidden=d_hidden,
                    n_heads_msa=n_heads_msa,
                    msa_head_width=msa_head_width,
                    is_final_block=(i == n_layers - 1),
                )
                for i in range(n_layers)
            ]
        )

    def set_chunk_size(self, chunk_size: int | None) -> None:
        for block in self.blocks:
            cast(MSAEncoderBlock, block).set_chunk_size(chunk_size)

    def forward(
        self,
        x_pair: Tensor,
        x_inputs: Tensor,
        msa_oh: Tensor,
        has_deletion: Tensor,
        deletion_value: Tensor,
        msa_attention_mask: Tensor,
    ) -> Tensor:
        # Every input tensor is pre-transposed to shape (b, l, m, ...) before this call.
        m_feat = torch.cat(
            [msa_oh, has_deletion.unsqueeze(-1), deletion_value.unsqueeze(-1)], dim=-1
        )
        m = self.embed(m_feat) + self.project_inputs(x_inputs).unsqueeze(2)
        tok_mask = msa_attention_mask[:, :, 0].bool()
        pair_attention_mask = tok_mask.unsqueeze(2) & tok_mask.unsqueeze(1)
        for block in self.blocks:
            m, x_pair = block(m, x_pair, msa_attention_mask, pair_attention_mask)
        return x_pair