# Copyright 2026 Biohub. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
"""Shared building blocks for ESMFold2 HuggingFace model variants."""

from __future__ import annotations

import importlib
from functools import partial
from importlib.util import find_spec
from typing import ClassVar, cast

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

from .configuration_esmfold2 import ESMFold2Config
from .reproducibility import seed_context

_seed_context = seed_context

try:
    if find_spec("cuequivariance_ops_torch") is None:
        raise ImportError("cuequivariance_ops_torch is unavailable")
    cue_module = importlib.import_module("cuequivariance_torch")
    _cue_attn_pair_bias = cue_module.attention_pair_bias
    _cue_tri_mul = cue_module.triangle_multiplicative_update

    CUE_AVAILABLE = True
except (AttributeError, ImportError):
    _cue_attn_pair_bias = None  # type: ignore[assignment]
    _cue_tri_mul = None  # type: ignore[assignment]
    CUE_AVAILABLE = False

# Biohub ships optional source-built Triton helpers. FastPLMs does not bundle or
# compile them; these placeholders retain checkpoint-compatible control flow
# while the portable PyTorch path remains flat and self-contained.
_fused_pair_bias = None
_fused_trimul_with_residual = None
_FusedLNLinearSwiGLU = None
_FusedDropoutResidual = None
TRITON_KERNELS_AVAILABLE = False

BACKEND_FUSED = "fused"
BACKEND_CUEQ = "cuequivariance"
_VALID_BACKENDS = (None, BACKEND_FUSED, BACKEND_CUEQ)
MSA_CONDITIONING_INPUT_NAMES = (
    "msa",
    "msa_attention_mask",
    "has_deletion",
    "deletion_value",
    "deletion_mean",
)
PREPARED_AUXILIARY_INPUT_NAMES = (
    "pocket_feature",
    "gt_coords",
    "is_resolved",
    "frames_idx",
    "disto_cond",
    "disto_cond_mask",
)


def validate_kernel_backend(backend: str | None) -> None:
    """Fail before mutating modules when a named kernel cannot execute."""

    if backend not in _VALID_BACKENDS:
        raise ValueError(f"backend must be one of {_VALID_BACKENDS}, got {backend!r}")
    if backend == BACKEND_FUSED and not TRITON_KERNELS_AVAILABLE:
        raise RuntimeError(
            "backend='fused' is unavailable because FastPLMs does not bundle the "
            "source-built ESMFold2 Triton kernels."
        )
    if backend == BACKEND_CUEQ and not CUE_AVAILABLE:
        raise RuntimeError(
            "backend='cuequivariance' requires cuequivariance_torch and the CUDA 13 "
            "cuequivariance_ops_torch runtime. Install FastPLMs with the "
            "'structure,cueq' extras on a supported Linux CUDA 13 host."
        )


def validate_msa_conditioning_inputs(
    config: ESMFold2Config,
    *,
    msa: Tensor | None,
    msa_attention_mask: Tensor | None,
    has_deletion: Tensor | None,
    deletion_value: Tensor | None,
    deletion_mean: Tensor | None,
) -> None:
    """Reject MSA-derived tensors for checkpoints trained without MSA conditioning."""

    if config.msa_conditioning:
        return
    values = {
        "msa": msa,
        "msa_attention_mask": msa_attention_mask,
        "has_deletion": has_deletion,
        "deletion_value": deletion_value,
        "deletion_mean": deletion_mean,
    }
    provided = sorted(name for name, value in values.items() if value is not None)
    if provided:
        raise ValueError(
            "This ESMFold2 checkpoint was trained without MSA conditioning and rejects "
            f"MSA-derived inputs: {', '.join(provided)}."
        )


def validate_prepared_auxiliary_inputs(
    *,
    pocket_feature: Tensor | None,
    disto_cond: Tensor | None,
    disto_cond_mask: Tensor | None,
) -> None:
    """Accept inert prepared fields and reject unsupported conditioning."""

    if pocket_feature is not None and torch.any(pocket_feature != 0).item():
        raise NotImplementedError(
            "The published ESMFold2 forward does not consume pocket conditioning; "
            "nonzero pocket_feature values are unsupported."
        )
    distogram_is_active = (disto_cond is not None and torch.any(disto_cond != 0).item()) or (
        disto_cond_mask is not None and torch.any(disto_cond_mask).item()
    )
    if distogram_is_active:
        raise NotImplementedError(
            "The published ESMFold2 forward does not consume distogram conditioning; "
            "nonzero disto_cond or disto_cond_mask values are unsupported."
        )


def _fused_active(module: nn.Module, tensor: Tensor) -> bool:
    """Return whether an optional fused implementation can handle this call."""
    return (
        TRITON_KERNELS_AVAILABLE
        and getattr(module, "_kernel_backend", None) == BACKEND_FUSED
        and not torch.is_grad_enabled()
        and tensor.is_cuda
    )


def _cueq_active(module: nn.Module) -> bool:
    return CUE_AVAILABLE and getattr(module, "_kernel_backend", None) == BACKEND_CUEQ


class DropoutResidual(nn.Module):
    """``residual + dropout(delta)`` with row/col-shared dropout.

    Same signature on both paths. ``use_fused_kernels=True`` + ``batch_dim=1``
    routes through ``FusedDropoutResidual`` (single-pass over pair tensor,
    in-place residual add). Falls back to unfused otherwise.
    """

    def __init__(self, r: float, batch_dim: int, use_fused_kernels: bool = False) -> None:
        super().__init__()
        if isinstance(batch_dim, bool) or batch_dim not in (1, 2):
            raise ValueError(f"batch_dim must be 1 or 2, got {batch_dim}")
        self._use_fused_kernels = (
            use_fused_kernels and batch_dim == 1 and _FusedDropoutResidual is not None
        )
        self._batch_dim = batch_dim
        self._r = r
        if self._use_fused_kernels:
            assert _FusedDropoutResidual is not None
            self._impl: nn.Module = _FusedDropoutResidual(r)
        else:
            self._impl = nn.Dropout(r)

    def forward(self, residual: Tensor, delta: Tensor) -> Tensor:
        if self._use_fused_kernels:
            return self._impl(residual, delta)
        # The unfused path broadcasts a row/column-shared mask M with shape (1, ...).
        if self._r == 0.0 or not self.training:
            return residual + delta
        shape = list(delta.shape)
        shape[self._batch_dim] = 1
        mask = self._impl(delta.new_ones(shape))
        return residual + delta * mask


# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
CHAR_VOCAB_SIZE: int = 64
MAX_CHARS: int = 4
XYZ_DIMS: int = 3
MAX_ATOMIC_NUMBER: int = 128

# Input feature dim = 3 + 1 + 1 + 128 + 64*4 = 389
ATOM_FEATURE_DIM: int = XYZ_DIMS + 1 + 1 + MAX_ATOMIC_NUMBER + CHAR_VOCAB_SIZE * MAX_CHARS


NUM_RES_TYPES: int = 33

_EPS = 1e-5

# Default for the quadratic triangle, OPM, and pair-transition operations.
# It caps peak memory so l around 2,000 fits on an 80 GiB GPU. At l=1,438,
# chunk=128 uses roughly 76 GiB, while chunk=64 leaves headroom for the
# largest foldbench targets. Pass None to disable chunking; this is faster for
# short sequences but prone to out-of-memory errors beyond l around 600.
_DEFAULT_CHUNK_SIZE = 64


# ===========================================================================
# MSA inference-time diversity augmentations
# ===========================================================================


def maybe_subsample_msa(
    msa: Tensor,
    msa_attention_mask: Tensor | None,
    has_deletion: Tensor | None,
    deletion_value: Tensor | None,
    *,
    max_depth: int | None,
    enabled: bool,
) -> tuple[Tensor, Tensor | None, Tensor | None, Tensor | None]:
    if not enabled or max_depth is None:
        return msa, msa_attention_mask, has_deletion, deletion_value

    depth = msa.size(1)
    if depth <= 1 or depth <= max_depth:
        return msa, msa_attention_mask, has_deletion, deletion_value

    indices = torch.zeros(max_depth, dtype=torch.long, device=msa.device)
    indices[1:] = torch.randperm(depth - 1, device=msa.device)[: max_depth - 1] + 1
    indices = indices.sort().values

    msa = msa[:, indices]
    if msa_attention_mask is not None:
        msa_attention_mask = msa_attention_mask[:, indices]
    if has_deletion is not None:
        has_deletion = has_deletion[:, indices]
    if deletion_value is not None:
        deletion_value = deletion_value[:, indices]
    return msa, msa_attention_mask, has_deletion, deletion_value


def maybe_apply_msa_column_masking(
    msa_attention_mask: Tensor | None,
    rate: float,
) -> Tensor | None:
    if msa_attention_mask is None or rate <= 0.0 or msa_attention_mask.size(1) <= 1:
        return msa_attention_mask

    batch_size, _, length = msa_attention_mask.shape
    col_keep = torch.rand(batch_size, length, device=msa_attention_mask.device) >= rate
    col_keep = col_keep.unsqueeze(1).expand_as(msa_attention_mask).clone()
    col_keep[:, 0, :] = True
    return msa_attention_mask.bool() & col_keep


# ===========================================================================
# Atom-token utilities
# ===========================================================================


def gather_token_to_atom(token_features: Tensor, atom_to_token_idx: Tensor) -> Tensor:
    """Broadcast per-token features to per-atom features using gather.

    Args:
        token_features: X with shape (b, l, d).
        atom_to_token_idx: I with shape (b, a), int64.

    Returns:
        X with shape (b, a, d).
    """
    idx = atom_to_token_idx.unsqueeze(-1).expand(-1, -1, token_features.size(-1))
    return torch.gather(token_features, 1, idx)


def scatter_atom_to_token(
    atom_features: Tensor,
    atom_to_token_idx: Tensor,
    n_tokens: int,
    atom_mask: Tensor | None = None,
) -> Tensor:
    """Aggregate per-atom features to per-token features (mean).

    Args:
        atom_features: X with shape (b, a, d).
        atom_to_token_idx: I with shape (b, a), int64.
        n_tokens: Token count l.
        atom_mask: M with shape (b, a), Boolean.

    Returns:
        X with shape (b, l, d).
    """
    batch_size, n_atoms, d_model = atom_features.shape
    n_out = n_tokens
    idx = atom_to_token_idx
    if atom_mask is not None:
        idx = torch.where(atom_mask, atom_to_token_idx, n_tokens)
        n_out = n_tokens + 1
    idx_expanded = idx.unsqueeze(-1).expand(batch_size, n_atoms, d_model)
    out = torch.zeros(
        batch_size,
        n_out,
        d_model,
        device=atom_features.device,
        dtype=atom_features.dtype,
    )
    out.scatter_reduce_(1, idx_expanded, atom_features, reduce="mean", include_self=False)
    return out[:, :n_tokens, :]


def gather_rep_atom_coords(coords: Tensor, rep_atom_idx: Tensor) -> Tensor:
    """Gather representative atom coordinates for each token.

    Args:
        coords: X with shape (b, a, 3).
        rep_atom_idx: I with shape (b, l), int64.

    Returns:
        X with shape (b, l, 3).
    """
    idx = rep_atom_idx.unsqueeze(-1).expand(-1, -1, coords.size(-1))
    return torch.gather(coords, 1, idx)


def _compute_intra_token_idx(atom_to_token: Tensor) -> Tensor:
    """Compute local atom index within each token (vectorised).

    Atoms belonging to the same token are contiguous, so this computes a
    running count that resets at each token boundary.

    Args:
        atom_to_token: I with shape (b, a), mapping each atom to its token.

    Returns:
        Index tensor I with shape (b, a) and values from zero through
        ``max_atoms_per_token - 1``.
    """
    same_as_prev = F.pad(atom_to_token[:, 1:] == atom_to_token[:, :-1], (1, 0), value=False)
    ones = torch.ones_like(atom_to_token)
    cumsum = torch.cumsum(ones, dim=-1)
    group_start = cumsum.masked_fill(same_as_prev, 0)
    group_start = torch.cummax(group_start, dim=-1).values
    return cumsum - group_start


def _categorical_mean(logits: Tensor, start: float, end: float) -> Tensor:
    """Expected value of a categorical distribution over evenly-spaced bins.

    Equivalent to ``CategoricalMixture(logits, bins=logits.shape[-1], start, end).mean()``.

    Args:
        logits: Logit tensor X with shape (..., n_bins).
        start: left boundary
        end: right boundary

    Returns:
        Expected value tensor Y with shape (...).
    """
    n_bins = logits.shape[-1]
    edges = torch.linspace(start, end, n_bins + 1, device=logits.device, dtype=torch.float32)
    v_bins = (edges[:-1] + edges[1:]) / 2  # V_bin has shape (n_bins,).
    return (logits.float().softmax(-1) @ v_bins.unsqueeze(1)).squeeze(-1)


# ===========================================================================
# Feature preparation and language-model projection
# ===========================================================================


class RowAttentionPooling(nn.Module):
    """Row-wise attention pooling: attn_proj, out_proj."""

    def __init__(self, d_pair: int, d_single: int) -> None:
        super().__init__()
        self.attn_proj = nn.Linear(d_pair, 1, bias=False)
        self.out_proj = nn.Linear(d_pair, d_single, bias=False)

    def forward(self, z: Tensor, mask: Tensor) -> Tensor:
        scores = self.attn_proj(z).squeeze(-1)
        mask_bias = torch.where(
            mask[:, None, :].bool(),
            torch.zeros_like(scores),
            torch.full_like(scores, -1e9),
        )
        scores = scores + mask_bias
        weights = F.softmax(scores, dim=-1)
        pooled = torch.einsum("bnm,bnmd->bnd", weights, z)
        return self.out_proj(pooled)


# ===========================================================================
# InputsEmbedder
# ===========================================================================


class InputsEmbedder(nn.Module):
    """Embeds input features including atom-level encoding via SWA attention."""

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

        self.atom_attention_encoder = ESMFold2AtomEncoder(
            d_atom=swa_cfg.d_atom,
            d_token=swa_cfg.d_token,
            n_blocks=swa_cfg.n_blocks,
            n_heads=swa_cfg.n_heads,
            swa_window_size=swa_cfg.swa_window_size,
            expansion_ratio=swa_cfg.expansion_ratio,
            structure_prediction=False,  # no coords_linear
            spatial_rope_base_frequency=swa_cfg.spatial_rope_base_frequency,
            n_spatial_rope_pairs_per_axis=swa_cfg.n_spatial_rope_pairs_per_axis,
            n_uid_rope_pairs=swa_cfg.n_uid_rope_pairs,
            uid_rope_base_frequency=swa_cfg.uid_rope_base_frequency,
        )

    def forward(
        self,
        aatype: Tensor,
        profile: Tensor,
        deletion_mean: Tensor,
        ref_pos: Tensor,
        atom_attention_mask: Tensor,
        ref_space_uid: Tensor,
        ref_charge: Tensor,
        ref_element: Tensor,
        ref_atom_name_chars: Tensor,
        atom_to_token: Tensor,
    ) -> Tensor:
        """Embed inputs into per-token features.

        Returns:
            X with shape (b, l, d_inputs), concatenating atom encoding,
            aatype, profile, and deletion mean.
        """
        a, _q, _c, _attn_params, _intermediates = self.atom_attention_encoder(
            ref_pos=ref_pos,
            atom_attention_mask=atom_attention_mask,
            ref_space_uid=ref_space_uid,
            ref_charge=ref_charge,
            ref_element=ref_element,
            ref_atom_name_chars=ref_atom_name_chars,
            atom_to_token=atom_to_token,
        )
        return torch.cat([a, aatype, profile, deletion_mean.unsqueeze(-1)], dim=-1)


# ===========================================================================
# ResIdxAsymIdSymIdEntityIdEncoding (trunk relative position)
# ===========================================================================


class ResIdxAsymIdSymIdEntityIdEncoding(nn.Module):
    """Embedding weight W has shape (d_pair, n_features).

    Here ``n_features = 2 * (2 * r_bins + 2) + 1 + (2 * c_bins + 2)``.

    For default r_bins=32, c_bins=2: 2*66 + 1 + 6 = 139.
    """

    def __init__(
        self,
        n_relative_residx_bins: int = 32,
        n_relative_chain_bins: int = 2,
        d_pair: int = 256,
    ) -> None:
        super().__init__()
        self.n_relative_residx_bins = n_relative_residx_bins
        self.n_relative_chain_bins = n_relative_chain_bins
        self.d_pair = d_pair

        n_feats_residue = 2 * n_relative_residx_bins + 2
        n_feats_token = 2 * n_relative_residx_bins + 2
        n_feats_chain = 2 * n_relative_chain_bins + 2
        n_feats_same_entity = 1
        total_feats = n_feats_residue + n_feats_token + n_feats_chain + n_feats_same_entity
        self.embed = nn.Linear(total_feats, d_pair, bias=False)

    def forward(
        self,
        residue_index: Tensor,
        asym_id: Tensor,
        sym_id: Tensor,
        entity_id: Tensor,
        token_index: Tensor,
    ) -> Tensor:
        bij_same_chain = asym_id.unsqueeze(2) == asym_id.unsqueeze(1)
        bij_same_residue = residue_index.unsqueeze(2) == residue_index.unsqueeze(1)
        bij_same_entity = entity_id.unsqueeze(2) == entity_id.unsqueeze(1)

        dij_residue = residue_index.unsqueeze(2) - residue_index.unsqueeze(1)
        dij_residue = torch.clip(
            dij_residue + self.n_relative_residx_bins,
            0,
            2 * self.n_relative_residx_bins,
        )
        dij_residue = torch.where(bij_same_chain, dij_residue, 2 * self.n_relative_residx_bins + 1)
        aij_rel_pos = F.one_hot(dij_residue, 2 * self.n_relative_residx_bins + 2)

        dij_token = torch.clip(
            token_index.unsqueeze(2) - token_index.unsqueeze(1) + self.n_relative_residx_bins,
            0,
            2 * self.n_relative_residx_bins,
        )
        dij_token = torch.where(
            bij_same_chain & bij_same_residue,
            dij_token,
            2 * self.n_relative_residx_bins + 1,
        )
        aij_rel_token = F.one_hot(dij_token, 2 * self.n_relative_residx_bins + 2)

        dij_chain = torch.clip(
            sym_id.unsqueeze(2) - sym_id.unsqueeze(1) + self.n_relative_chain_bins,
            0,
            2 * self.n_relative_chain_bins,
        )
        dij_chain = torch.where(bij_same_chain, 2 * self.n_relative_chain_bins + 1, dij_chain)
        aij_rel_chain = F.one_hot(dij_chain, 2 * self.n_relative_chain_bins + 2)

        feats = torch.cat(
            [
                aij_rel_pos.float(),
                aij_rel_token.float(),
                bij_same_entity.float().unsqueeze(-1),
                aij_rel_chain.float(),
            ],
            dim=-1,
        )

        return self.embed(feats)


# ===========================================================================
# SingleToPair (for LanguageModelShim)
# ===========================================================================


class SingleToPair(nn.Module):
    """downproject, output_mlp (Sequential of Linear, GELU, Linear)."""

    def __init__(self, input_dim: int, downproject_dim: int, output_dim: int) -> None:
        super().__init__()
        self.downproject = nn.Linear(input_dim, downproject_dim)
        self.output_mlp = nn.Sequential(
            nn.Linear(2 * downproject_dim, output_dim),
            nn.GELU(),
            nn.Linear(output_dim, output_dim),
        )

    def forward(self, x: Tensor) -> Tensor:
        x = self.downproject(x)
        x = torch.cat(
            [(x.unsqueeze(2) * x.unsqueeze(1)), (x.unsqueeze(2) - x.unsqueeze(1))],
            dim=3,
        )
        return self.output_mlp(x)


# ===========================================================================
# LanguageModelShim
# ===========================================================================


class LanguageModelShim(nn.Module):
    """Shim holding the trainable projection weights for LM integration.

    Contains:
    - base_z_combine: nn.Parameter with shape ``(n_layers + 1,)``
    - base_z_linear: Sequential(LayerNorm(d_model), Linear(d_model, d_z, bias=False))
    - base_z_mlp: Sequential(SingleToPair(d_z, d_z, d_z), LayerNorm(d_z))
    """

    def __init__(self, d_z: int = 256, d_model: int = 2560, num_layers: int = 80) -> None:
        super().__init__()

        self.base_z_mlp = nn.Sequential(SingleToPair(d_z, d_z, d_z), nn.LayerNorm(d_z))
        self.base_z_linear = nn.Sequential(
            nn.LayerNorm(d_model), nn.Linear(d_model, d_z, bias=False)
        )
        self.base_z_combine = nn.Parameter(torch.zeros(num_layers + 1))

    def project_sequence(
        self,
        hidden_states: Tensor,
        residue_mask: Tensor | None = None,
    ) -> Tensor:
        """Project all ESMC layer states into the learned sequence summary.

        Args:
            hidden_states: H with shape ``(b, l, n_layers + 1, d_model)``.
            residue_mask: Optional M with shape ``(b, l)``. Non-residue rows
                are set to zero in the returned tensor.

        Returns:
            Z with shape ``(b, l, d_z)``.
        """

        if hidden_states.ndim != 4:
            raise ValueError(
                "H must have shape (b, l, n_layers + 1, d_model), "
                f"got {tuple(hidden_states.shape)}."
            )
        expected_layers = self.base_z_combine.numel()
        if hidden_states.shape[-2] != expected_layers:
            raise ValueError(
                f"H contains {hidden_states.shape[-2]} states; expected "
                f"{expected_layers} in the official ESMC ordering."
            )
        expected_width = cast(nn.LayerNorm, self.base_z_linear[0]).normalized_shape[0]
        if hidden_states.shape[-1] != expected_width:
            raise ValueError(f"H has width {hidden_states.shape[-1]}; expected {expected_width}.")

        # H can be FP32 even when the folding checkpoint is loaded in BF16.
        # Match the learned projection parameters at this explicit boundary;
        # this preserves the official BF16 path and leaves FP32 models exact.
        projection_dtype = cast(nn.LayerNorm, self.base_z_linear[0]).weight.dtype
        hidden_states = hidden_states.to(dtype=projection_dtype)
        projected_states = self.base_z_linear(hidden_states)
        layer_weights = self.base_z_combine.softmax(dim=0)
        # Preserve Biohub's matmul path exactly so checkpoint inference does
        # not change through a different reduction order.
        projected = layer_weights @ projected_states
        if residue_mask is not None:
            if residue_mask.shape != hidden_states.shape[:2]:
                raise ValueError(
                    "M must have shape (b, l), got "
                    f"{tuple(residue_mask.shape)} for H {tuple(hidden_states.shape)}."
                )
            projected = projected * residue_mask.to(
                device=projected.device, dtype=projected.dtype
            ).unsqueeze(-1)
        return projected

    def forward(self, hidden_states: Tensor, *, lm_dropout: float = 0.0) -> Tensor:
        """Project pre-computed ESMC hidden states to pair representation.

        Args:
            hidden_states: H with shape ``(b, l, n_layers + 1, d_model)``.
            lm_dropout: Dropout probability applied to the pair
                representation after ``base_z_mlp``.

        Returns:
            Z_pair with shape ``(b, l, l, d_pair)``.
        """
        lm_z = self.project_sequence(hidden_states)
        lm_z = self.base_z_mlp(lm_z)
        if lm_dropout > 0:
            lm_z = F.dropout(lm_z, p=lm_dropout, training=True)
        return lm_z


# ===========================================================================
# ESMFold2ExperimentalModel: the top-level PreTrainedModel
# ===========================================================================


def compute_lm_hidden_states(
    esmc: nn.Module,
    input_ids: Tensor,
    asym_id: Tensor,
    residue_index: Tensor,
    mol_type: Tensor,
    token_mask: Tensor,
    pad_to_multiple: int | None = None,
    lm_mask_pct: float = 0.0,
    mask_token_id: int = 32,
) -> Tensor:
    """Run ESMC and return H with shape ``(b, l, n_states, d_model)``.

    Atom-tokenized modified residues (HYP, MSE, ACE, NH2, ...) span multiple
    structure tokens but share a single ``(asym_id, residue_index)`` key:
    collapse them to one LM token per residue before running the LM (the LM
    was trained on per-residue inputs, not per-atom), then scatter the
    hidden states back to the per-token layout.
    """
    b_size, l_size = input_ids.shape
    device = input_ids.device
    protein_mask = (mol_type == 0) & token_mask

    lm_input_list = []
    lm_lengths = []
    # Per-batch maps from (original protein-token index) to (LM input position).
    expand_maps: list[Tensor] = []
    for batch_index in range(b_size):
        mask_b = protein_mask[batch_index]
        ids_b = input_ids[batch_index][mask_b]
        asym_b = asym_id[batch_index][mask_b]
        res_b = residue_index[batch_index][mask_b]

        # Collapse: keep first token per (asym_id, residue_index) key, in
        # input order. ``inverse`` maps each original protein-token to its
        # collapsed residue index.
        keys = torch.stack((asym_b, res_b), dim=1)
        unique_keys, inverse = torch.unique(keys, dim=0, return_inverse=True)
        n_unique = unique_keys.size(0)
        token_positions = torch.arange(keys.size(0), device=device, dtype=torch.long)
        first_pos = torch.full((n_unique,), keys.size(0), device=device, dtype=torch.long)
        first_pos.scatter_reduce_(0, inverse, token_positions, reduce="amin", include_self=True)
        ordered = torch.argsort(first_pos)
        first_pos_ordered = first_pos[ordered]
        ids_collapsed = ids_b[first_pos_ordered]
        asym_collapsed = asym_b[first_pos_ordered]
        remap = torch.empty_like(ordered)
        remap[ordered] = torch.arange(n_unique, device=device, dtype=torch.long)
        inverse_ordered = remap[inverse]

        chain_ids = asym_collapsed.unique(sorted=True)
        # [BOS] chain1 [EOS BOS] chain2 ... [EOS]
        parts: list[Tensor] = [torch.tensor([0], device=device, dtype=ids_b.dtype)]
        # Per-chain LM positions accumulate; track them for the expand map.
        per_token_lm_pos = torch.empty(n_unique, device=device, dtype=torch.long)
        cursor = 1  # position 0 is the leading BOS
        for i, cid in enumerate(chain_ids):
            in_chain = (asym_collapsed == cid).nonzero(as_tuple=True)[0]
            parts.append(ids_collapsed[in_chain])
            per_token_lm_pos[in_chain] = torch.arange(
                cursor, cursor + in_chain.shape[0], device=device, dtype=torch.long
            )
            cursor += in_chain.shape[0]
            if i < len(chain_ids) - 1:
                parts.append(torch.tensor([2, 0], device=device, dtype=ids_b.dtype))
                cursor += 2  # EOS + BOS
        parts.append(torch.tensor([2], device=device, dtype=ids_b.dtype))
        lm_seq = torch.cat(parts)
        lm_input_list.append(lm_seq)
        lm_lengths.append(lm_seq.shape[0])

        # Map each original protein-token position to its LM input position.
        prot_pos_b = mask_b.nonzero(as_tuple=True)[0]
        expand_map = torch.full((l_size,), -1, device=device, dtype=torch.long)
        expand_map[prot_pos_b] = per_token_lm_pos[inverse_ordered]
        expand_maps.append(expand_map)

    # Pad the language-model input to its longest sequence. FP8 callers round
    # l to a multiple of 16 for Transformer Engine kernels.
    max_len = max(lm_lengths)
    if pad_to_multiple is not None and pad_to_multiple > 1:
        max_len = ((max_len + pad_to_multiple - 1) // pad_to_multiple) * pad_to_multiple
    lm_input_ids = torch.full(
        (b_size, max_len),
        1,
        device=device,
        dtype=input_ids.dtype,  # PAD=1
    )
    for batch_index in range(b_size):
        lm_input_ids[batch_index, : lm_lengths[batch_index]] = lm_input_list[batch_index]

    # sequence_id for chain-aware attention; PAD tokens get -1 (no attention).
    sequence_id = (lm_input_ids == 0).cumsum(dim=1) - 1  # BOS=0
    sequence_id = sequence_id.masked_fill(lm_input_ids == 1, -1)  # PAD=1

    if lm_mask_pct > 0.0:
        special = (lm_input_ids == 0) | (lm_input_ids == 1) | (lm_input_ids == 2)
        do_mask = (torch.rand(lm_input_ids.shape, device=device) < lm_mask_pct) & ~special
        lm_input_ids = lm_input_ids.masked_fill(do_mask, mask_token_id)

    with torch.inference_mode():
        esmc_out = esmc(input_ids=lm_input_ids, sequence_id=sequence_id, output_hidden_states=True)

    hidden_stack = esmc_out.hidden_states
    n_states, _, _, d_model = hidden_stack.shape
    result = torch.zeros(b_size, l_size, n_states, d_model, device=device, dtype=hidden_stack.dtype)
    for batch_index in range(b_size):
        M_i = protein_mask[batch_index]
        positions = expand_maps[batch_index][M_i]
        gathered = hidden_stack[:, batch_index, positions, :].permute(1, 0, 2)
        result[batch_index, M_i.nonzero(as_tuple=True)[0]] = gathered

    return result.detach()


# ===========================================================================
# TriangleMultiplicativeUpdate
# ===========================================================================


class TriangleMultiplicativeBlock(nn.Module):
    """Triangle multiplicative update block with gated signal routing."""

    _FLOW_TO_EINSUM: ClassVar[dict[str, str]] = {
        "outgoing": "bikd,bjkd->bijd",
        "incoming": "bkid,bkjd->bijd",
    }
    _VALID_FLOWS = ("outgoing", "incoming")

    def __init__(self, input_channels: int, latent_channels: int, flow: str) -> None:
        super().__init__()
        if flow not in self._FLOW_TO_EINSUM:
            raise ValueError(f"Invalid flow={flow!r}. Expected one of {self._VALID_FLOWS}.")

        self.input_channels = input_channels
        self.latent_channels = latent_channels
        self.flow = flow
        self._einsum_equation = self._FLOW_TO_EINSUM[flow]
        self.norm_start = nn.LayerNorm(self.input_channels, eps=_EPS)
        self.norm_mix = nn.LayerNorm(self.latent_channels, eps=_EPS)
        self.proj_bundle = nn.Linear(self.input_channels, 4 * self.latent_channels, bias=False)
        self.proj_emit = nn.Linear(self.latent_channels, self.input_channels, bias=False)
        self.proj_gate = nn.Linear(self.input_channels, self.input_channels, bias=False)

        self._use_kernels: bool = False
        # Default chunked for memory on long sequences; tests override with
        # ``set_chunk_size(None)`` for the unchunked path under bit-exact bf16
        # parity checks.
        self._chunk_size: int | None = 64

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

    def split_kernel_weights(self) -> tuple[Tensor, Tensor]:
        return (
            self.proj_bundle.weight[: 2 * self.latent_channels, :],
            self.proj_bundle.weight[2 * self.latent_channels :, :],
        )

    def _kernel_flow_direction(self) -> str:
        return self.flow

    def _triangular_contract(self, left_stream: Tensor, right_stream: Tensor) -> Tensor:
        return torch.einsum(self._einsum_equation, left_stream, right_stream)

    def _triangular_contract_chunked(
        self, left_stream: Tensor, right_stream: Tensor, chunk_size: int
    ) -> Tensor:
        """Compute the triangular einsum in chunks along the output i-dimension."""
        length = left_stream.shape[1] if self.flow == "outgoing" else left_stream.shape[2]
        chunks = []
        for start in range(0, length, chunk_size):
            end = min(start + chunk_size, length)
            if self.flow == "outgoing":
                chunk = torch.einsum(self._einsum_equation, left_stream[:, start:end], right_stream)
            else:
                chunk = torch.einsum(
                    self._einsum_equation, left_stream[:, :, start:end], right_stream
                )
            chunks.append(chunk)
        return torch.cat(chunks, dim=1)

    def forward(self, pair_grid: Tensor, visibility: Tensor | None = None) -> Tensor:
        if visibility is None:
            visibility = pair_grid.new_ones(pair_grid.shape[:-1])

        if self._use_kernels:
            p_in_weight, g_in_weight = self.split_kernel_weights()
            return _cue_tri_mul(  # type: ignore[misc]
                pair_grid,
                direction=self._kernel_flow_direction(),
                mask=visibility,
                norm_in_weight=self.norm_start.weight,
                norm_in_bias=self.norm_start.bias,
                p_in_weight=p_in_weight,
                g_in_weight=g_in_weight,
                norm_out_weight=self.norm_mix.weight,
                norm_out_bias=self.norm_mix.bias,
                p_out_weight=self.proj_emit.weight,
                g_out_weight=self.proj_gate.weight,
                eps=_EPS,
            )

        normalized_grid = self.norm_start(pair_grid)
        bundled = self.proj_bundle(normalized_grid)
        signal, gate_logits = bundled.split(2 * self.latent_channels, dim=-1)
        routed = signal * torch.sigmoid(gate_logits)
        routed = routed * visibility.unsqueeze(-1)

        left_stream, right_stream = routed.float().chunk(2, dim=-1)
        if self._chunk_size is not None:
            contracted = self._triangular_contract_chunked(
                left_stream, right_stream, self._chunk_size
            )
        else:
            contracted = self._triangular_contract(left_stream, right_stream)
        mixed = self.proj_emit(self.norm_mix(contracted))
        output_gate = torch.sigmoid(self.proj_gate(normalized_grid))
        return mixed * output_gate


class TriangleMultiplicativeUpdate(nn.Module):
    """Thin wrapper exposing the triangular mixer with explicit orientation (v3)."""

    def __init__(self, dim: int = 128, _outgoing: bool = True) -> None:
        super().__init__()
        flow = "outgoing" if _outgoing else "incoming"
        self._engine = TriangleMultiplicativeBlock(
            input_channels=dim, latent_channels=dim, flow=flow
        )

    def set_kernel_backend(self, backend: str | None) -> None:
        # Engine uses cueq when backend=="cuequivariance"; the "fused" backend
        # routes through the parent PairUpdateBlock's fused path (bypassing this).
        validate_kernel_backend(backend)
        self._engine._use_kernels = backend == BACKEND_CUEQ

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

    def forward(self, z: Tensor, mask: Tensor | None = None) -> Tensor:
        return self._engine(z, visibility=mask)


# ===========================================================================
# FoldingTrunk: Transition, PairUpdateBlock, FoldingTrunk
# ===========================================================================


class Transition(nn.Module):
    """LN + SwiGLU FFN with addmm-fused residual; optional Triton LN+w12+SwiGLU kernel."""

    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)
        # Default chunked; set_chunk_size(None) disables for bit-exact parity tests.
        self._chunk_size: int | None = 64
        self._fused_swiglu: nn.Module | None = None
        self._kernel_backend: str | None = None

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

    def set_kernel_backend(self, backend: str | None) -> None:
        """Install / uninstall FusedLNLinearSwiGLU (no cueq equivalent)."""
        validate_kernel_backend(backend)
        self._kernel_backend = backend
        if backend == BACKEND_FUSED and TRITON_KERNELS_AVAILABLE:
            assert _FusedLNLinearSwiGLU is not None
            d_model = self.norm.normalized_shape[0]
            d_inner = self.ffn.hidden_features
            has_ln_bias = self.norm.bias is not None
            device = self.ffn.w12.weight.device
            dtype = self.ffn.w12.weight.dtype
            fused = _FusedLNLinearSwiGLU(
                d_model=d_model,
                d_inner=d_inner,
                has_ln_bias=has_ln_bias,
                device=device,
                dtype=dtype,
            )
            with torch.no_grad():
                fused.LN_W.copy_(self.norm.weight)
                if has_ln_bias:
                    fused.LN_B.copy_(self.norm.bias)  # type: ignore[union-attr]
                # FusedLNLinearSwiGLU.W12 is (d_model, 2*d_inner); transpose nn.Linear once.
                fused.W12.copy_(self.ffn.w12.weight.t().contiguous())
            self._fused_swiglu = fused.eval().requires_grad_(False)
        else:
            self._fused_swiglu = None

    def _can_use_fused_path(self, x: Tensor) -> bool:
        return (
            _fused_active(self, x) and self._fused_swiglu is not None and x.dtype == torch.bfloat16
        )

    def _swiglu_pre_w3(self, x_normed: Tensor) -> Tensor:
        """SwiGLU through silu(x1)*x2, before the final w3."""
        ffn = self.ffn
        x12 = ffn.w12(x_normed)
        x1, x2 = x12.split(ffn.hidden_features, dim=-1)
        return F.silu(x1) * x2

    def _addmm_residual(self, x: Tensor, hidden: Tensor) -> Tensor:
        """x + w3(hidden) via single cuBLAS addmm: avoids transition-output allocation."""
        ffn = self.ffn
        x_shape = x.shape
        out = torch.addmm(
            x.contiguous().view(-1, x_shape[-1]),
            hidden.view(-1, hidden.shape[-1]),
            ffn.w3.weight.t(),
        )
        return out.view(x_shape)

    def forward(self, x: Tensor) -> Tensor:
        # Inference-only fast path (addmm-fused residual + pre-alloc out)
        #: diverges bit-exactly from ``x + ffn(norm(x))`` so we only use
        # it when grad is disabled (binder-design / bit-exact tests run
        # with grad on and need the reference path).
        if not torch.is_grad_enabled() and self._can_use_fused_path(x):
            fused = self._fused_swiglu
            assert fused is not None
            pre_w3 = fused
            if self._chunk_size is None or x.shape[1] <= self._chunk_size:
                hidden = pre_w3(x)
                return self._addmm_residual(x, hidden)
            out = torch.empty_like(x)
            for s in range(0, x.shape[1], self._chunk_size):
                e = min(s + self._chunk_size, x.shape[1])
                sl = x[:, s:e]
                hidden = pre_w3(sl)
                out[:, s:e] = self._addmm_residual(sl, hidden)
            return out
        # Reference path: bit-exact with main: x + ffn(norm(x)).
        if self._chunk_size is None or x.shape[1] <= self._chunk_size:
            return x + self.ffn(self.norm(x))
        out_list: 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_list.append(sl + self.ffn(self.norm(sl)))
        return torch.cat(out_list, dim=1)


class PairUpdateBlock(nn.Module):
    """tri_mul_out, tri_mul_in, pair_transition."""

    def __init__(self, d_pair: int = 256, expansion_ratio: int = 4) -> None:
        super().__init__()
        self.tri_mul_out = TriangleMultiplicativeUpdate(dim=d_pair, _outgoing=True)
        self.tri_mul_in = TriangleMultiplicativeUpdate(dim=d_pair, _outgoing=False)
        self.pair_transition = Transition(d_pair, expansion_ratio=expansion_ratio)
        self._kernel_backend: str | None = None
        # Row-shared dropout-residual; r=0 for inference (HF model is inference-only).
        # backend='fused' swaps in the FusedDropoutResidual Triton kernel.
        self.row_drop = DropoutResidual(0.0, batch_dim=1, use_fused_kernels=False)

    def set_kernel_backend(self, backend: str | None) -> None:
        if backend not in _VALID_BACKENDS:
            raise ValueError(f"backend must be one of {_VALID_BACKENDS}, got {backend!r}")
        self.tri_mul_out.set_kernel_backend(backend)
        self.tri_mul_in.set_kernel_backend(backend)
        self.pair_transition.set_kernel_backend(backend)
        self._kernel_backend = backend
        self.row_drop = DropoutResidual(
            0.0, batch_dim=1, use_fused_kernels=(backend == BACKEND_FUSED)
        )

    def set_chunk_size(self, chunk_size: int | None) -> None:
        self.tri_mul_out.set_chunk_size(chunk_size)
        self.tri_mul_in.set_chunk_size(chunk_size)
        self.pair_transition.set_chunk_size(chunk_size)

    def _can_use_fused_trimul_with_residual(self, pair: Tensor) -> bool:
        return _fused_active(self, pair) and pair.dtype == torch.bfloat16

    def _fused_trimul_with_residual(
        self, pair: Tensor, direction: str, pair_attention_mask: Tensor | None
    ) -> Tensor:
        """Fused TriMul+residual call; weights from the corresponding engine."""
        tri = self.tri_mul_out if direction == "outgoing" else self.tri_mul_in
        engine: TriangleMultiplicativeBlock = tri._engine  # type: ignore[assignment]
        p_in_weight, g_in_weight = engine.split_kernel_weights()

        def _bf16(t: Tensor) -> Tensor:
            return t if t.dtype == torch.bfloat16 else t.to(torch.bfloat16)

        return _fused_trimul_with_residual(  # type: ignore[misc]
            pair,
            direction,
            residual=pair,
            drop_mask=None,  # inference: no dropout, matches internal's eval path
            norm_in_weight=_bf16(engine.norm_start.weight),
            norm_in_bias=_bf16(engine.norm_start.bias),
            p_in_weight=_bf16(p_in_weight),
            g_in_weight=_bf16(g_in_weight),
            norm_out_weight=_bf16(engine.norm_mix.weight),
            norm_out_bias=_bf16(engine.norm_mix.bias),
            p_out_weight=_bf16(engine.proj_emit.weight),
            g_out_weight=_bf16(engine.proj_gate.weight),
            mask=pair_attention_mask,
            eps=_EPS,
        )

    def forward(self, pair: Tensor, pair_attention_mask: Tensor | None = None) -> Tensor:
        if self._can_use_fused_trimul_with_residual(pair):
            pair = self._fused_trimul_with_residual(pair, "outgoing", pair_attention_mask)
            pair = self._fused_trimul_with_residual(pair, "incoming", pair_attention_mask)
        else:
            pair = self.row_drop(pair, self.tri_mul_out(pair, mask=pair_attention_mask))
            pair = self.row_drop(pair, self.tri_mul_in(pair, mask=pair_attention_mask))
        pair = self.pair_transition(pair)
        return pair


class FoldingTrunk(nn.Module):
    """ModuleList of PairUpdateBlocks."""

    def __init__(self, n_layers: int = 24, d_pair: int = 256, expansion_ratio: int = 4) -> None:
        super().__init__()
        self.blocks = nn.ModuleList(
            [
                PairUpdateBlock(d_pair=d_pair, expansion_ratio=expansion_ratio)
                for _ in range(n_layers)
            ]
        )

    def set_kernel_backend(self, backend: str | None) -> None:
        for block in self.blocks:
            cast(PairUpdateBlock, block).set_kernel_backend(backend)

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

    def forward(self, pair: Tensor, pair_attention_mask: Tensor | None = None) -> Tensor:
        # Cast the pair tensor to BF16 when the fused triangle backend is enabled
        # (its bwd kernel requires bf16). Other backends keep the input dtype.
        orig_dtype = pair.dtype
        fused_on = (
            len(self.blocks) > 0
            and getattr(self.blocks[0], "_kernel_backend", None) == BACKEND_FUSED
        )
        if pair.is_cuda and fused_on and orig_dtype != torch.bfloat16:
            pair = pair.to(torch.bfloat16)
        for block in self.blocks:
            fn = partial(block, pair_attention_mask=pair_attention_mask)
            if torch.is_grad_enabled():
                pair = checkpoint(fn, pair, use_reentrant=False)  # pyright: ignore
            else:
                pair = fn(pair)
        if pair.dtype != orig_dtype:
            pair = pair.to(orig_dtype)
        return pair


# ===========================================================================
# MSA Encoder
# ===========================================================================


class OuterProductMean(nn.Module):
    """Outer-product mean: maps an MSA representation into a pair update.

    The order of the ``/ n_valid`` divide vs. the ``Wout`` projection is
    selectable via ``divide_outer_before_proj`` because different ESMFold2
    checkpoints were trained with different orderings:

    * ``False`` (default): ``Wout(outer) / n_valid``: the projection bias
      is scaled by 1/n_valid alongside the outer product.
    * ``True``: ``Wout(outer / n_valid)``: the projection bias is added
      unscaled, post-divide.
    """

    def __init__(
        self,
        d_msa: int,
        d_hidden: int,
        d_pair: int,
        divide_outer_before_proj: bool = False,
    ) -> None:
        super().__init__()
        self.d_hidden = d_hidden
        self.divide_outer_before_proj = divide_outer_before_proj
        self.norm = nn.LayerNorm(d_msa)
        self.W = nn.Linear(d_msa, 2 * d_hidden, bias=False)
        self.Wout = nn.Linear(d_hidden * d_hidden, d_pair, bias=True)
        # Off for bit-exact bf16; ``set_chunk_size(64)`` for long sequences.
        self._chunk_size: int | None = None

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

    def forward(self, m: Tensor, msa_attention_mask: Tensor) -> Tensor:
        m_norm = self.norm(m)
        x = self.W(m_norm) * msa_attention_mask.unsqueeze(-1).to(m_norm.dtype)
        a, b = x.chunk(2, dim=-1)
        mask_f = msa_attention_mask.to(a.dtype)
        n_valid = (mask_f @ mask_f.transpose(-1, -2)).unsqueeze(-1).clamp(min=1.0)
        if self._chunk_size is None:
            outer = torch.einsum("bimc,bjmd->bijcd", a, b).flatten(-2)
            if self.divide_outer_before_proj:
                return self.Wout(outer / n_valid)
            return self.Wout(outer) / n_valid
        # Chunk along the left (i) axis so the peak einsum intermediate is
        # X uses shape (b, chunk, l, c, d) instead of (b, l, l, c, d).
        length = a.shape[1]
        out_chunks: list[Tensor] = []
        for start in range(0, length, self._chunk_size):
            end = min(start + self._chunk_size, length)
            outer_chunk = torch.einsum("bimc,bjmd->bijcd", a[:, start:end], b).flatten(-2)
            if self.divide_outer_before_proj:
                out_chunks.append(self.Wout(outer_chunk / n_valid[:, start:end]))
            else:
                out_chunks.append(self.Wout(outer_chunk) / n_valid[:, start:end])
        return torch.cat(out_chunks, dim=1)


class MSAPairWeightedAveraging(nn.Module):
    """Pair-biased MSA row update (AF3 Supplement Algorithm 10)."""

    def __init__(self, d_msa: int, d_pair: int, n_heads: int = 8, head_width: int = 32) -> None:
        super().__init__()
        self.n_heads = n_heads
        self.head_width = head_width
        self.norm_single = nn.LayerNorm(d_msa)
        self.compute_bias = nn.Sequential(
            nn.LayerNorm(d_pair), nn.Linear(d_pair, n_heads, bias=False)
        )
        self.Wv = nn.Linear(d_msa, n_heads * head_width, bias=False)
        self.Wgate = nn.Linear(d_msa, n_heads * head_width, bias=False)
        self.Wout = nn.Linear(n_heads * head_width, d_msa, bias=False)

    def forward(self, msa_repr: Tensor, pair_repr: Tensor, pair_attention_mask: Tensor) -> Tensor:
        """
        Args:
            msa_repr: X with shape (b, l, m, d_msa).
            pair_repr: Z with shape (b, l, l, d_pair).
            pair_attention_mask: M with shape (b, l, l).
        Returns:
            X with shape (b, l, m, d_msa).
        """
        batch_size, length, depth, _ = msa_repr.shape
        n_heads, head_width = self.n_heads, self.head_width

        msa_normed = self.norm_single(msa_repr)
        bias = self.compute_bias(pair_repr)  # A has shape (b, l, l, n_heads).
        bias.masked_fill_(~pair_attention_mask.unsqueeze(-1).bool(), -1e5)
        attn = torch.softmax(bias, dim=-2)  # softmax over j

        v = self.Wv(msa_normed).reshape(batch_size, length, depth, n_heads, head_width)
        gate = torch.sigmoid(self.Wgate(msa_normed)).reshape(
            batch_size, length, depth, n_heads, head_width
        )

        output = torch.einsum("bijh,bjmhd,bimhd->bimhd", attn, v, gate)
        return self.Wout(output.reshape(batch_size, length, depth, n_heads * head_width))


# ===========================================================================
# Atom and diffusion stack
# ===========================================================================


class TransitionLayer(nn.Module):
    """SwiGLU transition: norm -> a_proj, b_proj -> silu(a)*b -> out_proj."""

    def __init__(self, d_model: int, n: int, eps: float = 1e-5) -> None:
        super().__init__()
        hidden = n * d_model
        self.norm = nn.LayerNorm(d_model, eps=eps)
        self.a_proj = nn.Linear(d_model, hidden, bias=False)
        self.b_proj = nn.Linear(d_model, hidden, bias=False)
        self.out_proj = nn.Linear(hidden, d_model, bias=False)

    def forward(self, x: Tensor) -> Tensor:
        x = self.norm(x)
        a = self.a_proj(x)
        b = self.b_proj(x)
        return self.out_proj(F.silu(a) * b)


# ===========================================================================
# AdaptiveLayerNorm (used in DiffusionTransformer)
# ===========================================================================


class AdaptiveLayerNorm(nn.Module):
    """Adaptive layer normalization (adaLN-Zero)."""

    def __init__(self, d_model: int, d_cond: int, eps: float = 1e-5) -> None:
        super().__init__()
        self.d_model = d_model
        self.d_cond = d_cond
        self.eps = eps
        self.s_scale = nn.Parameter(torch.ones(d_cond))
        self.s_gate = nn.Linear(d_cond, d_model, bias=True)
        self.s_shift = nn.Linear(d_cond, d_model, bias=False)

    def forward(self, a: Tensor, s: Tensor) -> Tensor:
        a_norm = F.layer_norm(a, (self.d_model,), None, None, self.eps)
        s_norm = F.layer_norm(s, (self.d_cond,), self.s_scale, None, self.eps)
        return torch.sigmoid(self.s_gate(s_norm)) * a_norm + self.s_shift(s_norm)


# ===========================================================================
# FourierEmbedding
# ===========================================================================


class FourierEmbedding(nn.Module):
    """Fourier embedding: cos(2*pi*(t*w + b))."""

    w: Tensor
    b: Tensor

    def __init__(self, c: int) -> None:
        super().__init__()
        self.c = c
        self.register_buffer("w", torch.randn(c))
        self.register_buffer("b", torch.randn(c))

    def forward(self, t_hat: Tensor) -> Tensor:
        t = torch.as_tensor(t_hat, device=self.w.device, dtype=self.w.dtype).reshape(-1)
        return torch.cos(2.0 * torch.pi * (t[:, None] * self.w[None, :] + self.b[None, :]))


# ===========================================================================
# SwiGLU / SwiGLUMLP
# ===========================================================================


def _compute_swiglu_hidden_size(d_model: int, expansion_ratio: int) -> int:
    return expansion_ratio * d_model


class SwiGLU(nn.Module):
    """SwiGLU with packed w12 and output w3."""

    def __init__(
        self,
        in_features: int,
        hidden_features: int,
        out_features: int | None = None,
        bias: bool = True,
    ) -> None:
        super().__init__()
        out_features = out_features or in_features
        self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias)
        self.w3 = nn.Linear(hidden_features, out_features, bias=bias)
        self.hidden_features = hidden_features

    def forward(self, x: Tensor) -> Tensor:
        x12 = self.w12(x)
        x1, x2 = x12.split(self.hidden_features, dim=-1)
        hidden = F.silu(x1) * x2
        return self.w3(hidden)


class SwiGLUMLP(SwiGLU):
    """SwiGLU MLP with packed weights, no bias."""

    def __init__(self, d_model: int, expansion_ratio: int = 4, bias: bool = False) -> None:
        hidden = _compute_swiglu_hidden_size(d_model, expansion_ratio)
        super().__init__(
            in_features=d_model, hidden_features=hidden, out_features=d_model, bias=bias
        )


# ===========================================================================
# SWA Atom Attention components
# ===========================================================================


def _rotate_half(x: Tensor) -> Tensor:
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)


def apply_rotary_emb_3d(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor:
    """Apply RoPE with batch-dependent cos/sin.

    Args:
        x: X with shape (b, l, h, d).
        cos: C with shape (b, l, d / 2).
        sin: S with shape (b, l, d / 2).
    """
    ro_dim = cos.shape[-1] * 2
    cos = cos.unsqueeze(2).repeat(1, 1, 1, 2)
    sin = sin.unsqueeze(2).repeat(1, 1, 1, 2)
    return torch.cat(
        [x[..., :ro_dim] * cos + _rotate_half(x[..., :ro_dim]) * sin, x[..., ro_dim:]],
        dim=-1,
    )


@torch.compiler.disable
def build_3d_rope(
    ref_pos: Tensor,
    ref_space_uid: Tensor,
    head_dim: int,
    n_spatial_per_axis: int = 4,
    n_uid_pairs: int = 2,
    spatial_base_freq: float = 10000.0,
    uid_base_freq: float = 10.0,
) -> tuple[Tensor, Tensor]:
    """Build cos/sin for 3D RoPE + UID RoPE."""
    device = ref_pos.device
    batch_size, n_atoms = ref_pos.shape[:2]
    half_dim = head_dim // 2
    n_spatial_total = 3 * n_spatial_per_axis

    spatial_inv_freq = 1.0 / (
        spatial_base_freq
        ** (
            torch.arange(0, n_spatial_per_axis, dtype=torch.float32, device=device)
            / n_spatial_per_axis
        )
    )
    uid_inv_freq = 1.0 / (
        uid_base_freq
        ** (torch.arange(0, n_uid_pairs, dtype=torch.float32, device=device) / n_uid_pairs)
    )

    pos_f32 = ref_pos.float()
    spatial_freqs = torch.einsum("bna,k->bnak", pos_f32, spatial_inv_freq)
    spatial_freqs = spatial_freqs.reshape(batch_size, n_atoms, n_spatial_total)

    uid_f32 = ref_space_uid.float()
    uid_freqs = torch.einsum("bn,k->bnk", uid_f32, uid_inv_freq)

    n_active = n_spatial_total + n_uid_pairs
    freqs = torch.cat([spatial_freqs, uid_freqs], dim=-1)

    if n_active < half_dim:
        padding = torch.zeros(
            batch_size,
            n_atoms,
            half_dim - n_active,
            device=device,
            dtype=torch.float32,
        )
        freqs = torch.cat([freqs, padding], dim=-1)

    cos = freqs.cos().to(torch.bfloat16)
    sin = freqs.sin().to(torch.bfloat16)
    return cos, sin


def qk_norm(x: Tensor) -> Tensor:
    return F.rms_norm(x, (x.size(-1),)).to(x.dtype)


# ===========================================================================
# SwiGLUFFN (atom transformer blocks)
# ===========================================================================


class SwiGLUFFN(nn.Module):
    """SwiGLU FFN with rounded hidden size for hardware alignment."""

    def __init__(self, d_model: int, expansion_ratio: int = 2) -> None:
        super().__init__()
        hidden_size = ((expansion_ratio * (d_model // 3) * 2) + 255) // 256 * 256
        self.w_up = nn.Linear(d_model, 2 * hidden_size, bias=False)
        self.w_down = nn.Linear(hidden_size, d_model, bias=False)

    def forward(self, x: Tensor) -> Tensor:
        x = x.to(self.w_up.weight.dtype)
        x1, x2 = self.w_up(x).chunk(2, dim=-1)
        return self.w_down(F.silu(x1) * x2)


# ===========================================================================
# SWA3DRoPEAttention
# ===========================================================================


class SWA3DRoPEAttention(nn.Module):
    """Sliding window attention with 3D RoPE. Has Wqkv, gate_proj, out_proj."""

    def __init__(self, d_model: int, n_heads: int, half_window: int = 64) -> None:
        super().__init__()
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads
        self.scale = self.head_dim**-0.5
        self.half_window = half_window

        self.Wqkv = nn.Linear(d_model, 3 * d_model, bias=False)
        self.out_proj = nn.Linear(d_model, d_model, bias=False)
        self.gate_proj = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x: Tensor, attention_params: tuple) -> Tensor:
        batch_size, n_atoms = x.shape[:2]
        cos, sin = attention_params[0], attention_params[1]

        x_input = x
        qkv = self.Wqkv(x)
        qkv = qkv.view(batch_size, n_atoms, 3, self.n_heads, self.head_dim).permute(2, 0, 1, 3, 4)
        q, k, v = qkv.unbind(0)
        q, k = qk_norm(q), qk_norm(k)

        q = apply_rotary_emb_3d(q, cos, sin)
        k = apply_rotary_emb_3d(k, cos, sin)

        input_dtype = q.dtype
        if q.dtype not in (torch.float16, torch.bfloat16):
            q, k, v = q.bfloat16(), k.bfloat16(), v.bfloat16()

        # ESMFold2 does not advertise FlashAttention. Keep this atom path on
        # PyTorch. Models that advertise FlashAttention dispatch through the
        # precompiled Hugging Face kernels interface in fastplms.attention.
        q_t = q.transpose(1, 2)
        k_t = k.transpose(1, 2)
        v_t = v.transpose(1, 2)
        attn = torch.matmul(q_t, k_t.transpose(-2, -1)) * self.scale
        attn = F.softmax(attn, dim=-1)
        out = torch.matmul(attn, v_t).transpose(1, 2)

        out = out.to(input_dtype).reshape(  # type: ignore[union-attr]
            batch_size, n_atoms, -1
        )
        out = out * torch.sigmoid(self.gate_proj(x_input))
        return self.out_proj(out)


# ===========================================================================
# SWAAtomBlock, SWAAtomTransformer
# ===========================================================================


def _rms_adaln_raw(x: Tensor, scale: Tensor, shift: Tensor) -> Tensor:
    return F.rms_norm(x, (x.shape[-1],)) * (1 + scale) + shift


def _gated_residual_raw(x: Tensor, gate: Tensor, y: Tensor) -> Tensor:
    return x + gate * y


class SWAAtomBlock(nn.Module):
    """adaLN-Zero + SWA attention + SwiGLU FFN.

    Creates adaln_modulation = Sequential(SiLU(), Linear) -> keys like adaln_modulation.1.weight
    """

    def __init__(
        self,
        d_atom: int,
        n_heads: int,
        half_window: int = 64,
        expansion_ratio: int = 2,
        use_compile_fusions: bool = False,
    ) -> None:
        super().__init__()
        self.attn_norm = nn.RMSNorm(d_atom, elementwise_affine=False)
        self.ffn_norm = nn.RMSNorm(d_atom, elementwise_affine=False)

        adaln_linear = nn.Linear(d_atom, 6 * d_atom, bias=False)
        nn.init.zeros_(adaln_linear.weight)
        self.adaln_modulation = nn.Sequential(nn.SiLU(), adaln_linear)

        self.attn = SWA3DRoPEAttention(d_atom, n_heads, half_window=half_window)
        self.ffn = SwiGLUFFN(d_atom, expansion_ratio)

        self._rms_adaln = torch.compile(_rms_adaln_raw) if use_compile_fusions else _rms_adaln_raw
        self._gated_residual = (
            torch.compile(_gated_residual_raw) if use_compile_fusions else _gated_residual_raw
        )

    def forward(self, x: Tensor, c_l: Tensor, attention_params: tuple) -> Tensor:
        mod = self.adaln_modulation(c_l)
        if mod.dim() == 2:
            mod = mod.unsqueeze(1)
        shift_a, scale_a, gate_a, shift_f, scale_f, gate_f = mod.chunk(6, dim=-1)

        attn_input = self._rms_adaln(x, scale_a, shift_a)
        attn_out = self.attn(attn_input, attention_params)
        x = self._gated_residual(x, gate_a, attn_out)

        ffn_input = self._rms_adaln(x, scale_f, shift_f)
        ffn_out = self.ffn(ffn_input)
        x = self._gated_residual(x, gate_f, ffn_out)
        return x


class SWAAtomTransformer(nn.Module):
    """Stack of SWAAtomBlocks."""

    def __init__(
        self,
        d_atom: int = 128,
        n_blocks: int = 3,
        n_heads: int = 4,
        swa_window_size: int = 128,
        expansion_ratio: int = 2,
        spatial_rope_base_frequency: float = 20.0,
        n_spatial_rope_pairs_per_axis: int = 2,
        n_uid_rope_pairs: int = 10,
        uid_rope_base_frequency: float = 10000.0,
    ) -> None:
        super().__init__()
        self.swa_window_size = swa_window_size
        self.head_dim = d_atom // n_heads
        self.spatial_rope_base_frequency = spatial_rope_base_frequency
        self.n_spatial_rope_pairs_per_axis = n_spatial_rope_pairs_per_axis
        self.n_uid_rope_pairs = n_uid_rope_pairs
        self.uid_rope_base_frequency = uid_rope_base_frequency

        self.blocks = nn.ModuleList(
            [
                SWAAtomBlock(
                    d_atom=d_atom,
                    n_heads=n_heads,
                    half_window=swa_window_size // 2,
                    expansion_ratio=expansion_ratio,
                )
                for _ in range(n_blocks)
            ]
        )

    def _build_3d_rope(self, ref_pos: Tensor, ref_space_uid: Tensor) -> tuple[Tensor, Tensor]:
        return build_3d_rope(
            ref_pos=ref_pos,
            ref_space_uid=ref_space_uid,
            head_dim=self.head_dim,
            n_spatial_per_axis=self.n_spatial_rope_pairs_per_axis,
            n_uid_pairs=self.n_uid_rope_pairs,
            spatial_base_freq=self.spatial_rope_base_frequency,
            uid_base_freq=self.uid_rope_base_frequency,
        )

    def forward(
        self,
        q_l: Tensor,
        c_l: Tensor,
        attention_params: tuple,
        return_intermediates: bool = False,
    ) -> Tensor | tuple[Tensor, list[Tensor]]:
        intermediates: list[Tensor] = []
        for block in self.blocks:
            q_l = block(q_l, c_l, attention_params)
            if return_intermediates:
                intermediates.append(q_l)
        if return_intermediates:
            return q_l, intermediates
        return q_l


# ===========================================================================
# ESMFold2AtomEncoder (for both inputs_embedder and diffusion_module)
# ===========================================================================


@torch.compiler.disable
def _prepare_atom_encoder_metadata(
    atom_attention_mask: Tensor,
    atom_to_token: Tensor,
    num_diffusion_samples: int,
) -> tuple[Tensor, Tensor, Tensor, int, int]:
    """Prepare mask-derived atom metadata outside compiled diffusion graphs."""
    mask_exp = atom_attention_mask.repeat_interleave(num_diffusion_samples, 0)
    seqlens = mask_exp.sum(dim=-1, dtype=torch.int32)
    indices = torch.nonzero(mask_exp.flatten(), as_tuple=False).flatten()
    max_seqlen = int(seqlens.max().item())
    cu_seqlens = F.pad(torch.cumsum(seqlens, dim=0, dtype=torch.int32), (1, 0))
    n_tokens = int(atom_to_token.max().item()) + 1
    return mask_exp, indices, cu_seqlens, max_seqlen, n_tokens


class ESMFold2AtomEncoder(nn.Module):
    """Encode atom inputs with normalization and sliding-window attention.

    Args:
        d_atom: atom hidden dim
        d_token: token dim for atom_to_token aggregation
        n_blocks, n_heads, swa_window_size, expansion_ratio: transformer params
        structure_prediction: if True, creates coords_linear and uses full d_token
        spatial_rope_base_frequency, n_spatial_rope_pairs_per_axis,
        n_uid_rope_pairs, uid_rope_base_frequency: 3D RoPE config
    """

    def __init__(
        self,
        d_atom: int = 128,
        d_token: int = 768,
        n_blocks: int = 3,
        n_heads: int = 4,
        swa_window_size: int = 128,
        expansion_ratio: int = 2,
        structure_prediction: bool = True,
        spatial_rope_base_frequency: float = 20.0,
        n_spatial_rope_pairs_per_axis: int = 2,
        n_uid_rope_pairs: int = 10,
        uid_rope_base_frequency: float = 10000.0,
    ) -> None:
        super().__init__()
        self.d_atom = d_atom
        self.d_token = d_token
        self.structure_prediction = structure_prediction

        self.atom_linear = nn.Linear(ATOM_FEATURE_DIM, d_atom, bias=False)
        self.atom_norm = nn.LayerNorm(d_atom)

        if structure_prediction:
            self.coords_linear = nn.Linear(6, d_atom, bias=False)

        self.atom_transformer = SWAAtomTransformer(
            d_atom=d_atom,
            n_blocks=n_blocks,
            n_heads=n_heads,
            swa_window_size=swa_window_size,
            expansion_ratio=expansion_ratio,
            spatial_rope_base_frequency=spatial_rope_base_frequency,
            n_spatial_rope_pairs_per_axis=n_spatial_rope_pairs_per_axis,
            n_uid_rope_pairs=n_uid_rope_pairs,
            uid_rope_base_frequency=uid_rope_base_frequency,
        )

        # Output aggregation: d_token for structure prediction, d_token//2 for inputs
        out_dim = d_token if structure_prediction else d_token // 2
        self.atom_to_token_linear = nn.Linear(d_atom, out_dim, bias=False)

    def forward(
        self,
        ref_pos: Tensor,
        atom_attention_mask: Tensor,
        ref_space_uid: Tensor,
        ref_charge: Tensor,
        ref_element: Tensor,
        ref_atom_name_chars: Tensor,
        atom_to_token: Tensor,
        r_l: Tensor | None = None,
        pred_r1: Tensor | None = None,
        s_i: Tensor | None = None,
        z_ij: Tensor | None = None,
        num_diffusion_samples: int = 1,
        return_intermediates: bool = False,
        inference_cache: dict | None = None,
    ) -> tuple[Tensor, Tensor, Tensor, tuple, list[Tensor]]:
        """Returns (a, q, c, attention_params, intermediates).

        ``inference_cache`` caches step-invariant tensors (c_base, 3D RoPE,
        attention indices, n_tokens) across diffusion steps.
        """
        batch_size, n_atoms = ref_pos.shape[:2]

        layer_cache = None
        if inference_cache is not None:
            layer_cache = inference_cache.setdefault("atomencoder", {})

        if layer_cache is None or len(layer_cache) == 0:
            atom_feats = torch.cat(
                [
                    ref_pos,
                    ref_charge.unsqueeze(-1),
                    atom_attention_mask.unsqueeze(-1),
                    ref_element,
                    ref_atom_name_chars.reshape(batch_size, n_atoms, MAX_CHARS * CHAR_VOCAB_SIZE),
                ],
                dim=-1,
            )
            c_base = self.atom_norm(self.atom_linear(atom_feats))
            cos, sin = self.atom_transformer._build_3d_rope(ref_pos, ref_space_uid)
            cos = cos.repeat_interleave(num_diffusion_samples, 0)
            sin = sin.repeat_interleave(num_diffusion_samples, 0)
            mask_exp, indices, cu_seqlens, max_seqlen, n_tokens = (
                _prepare_atom_encoder_metadata(
                    atom_attention_mask,
                    atom_to_token,
                    num_diffusion_samples,
                )
            )
            attention_params = (cos, sin, indices, cu_seqlens, max_seqlen)
            if layer_cache is not None:
                layer_cache["c_base"] = c_base
                layer_cache["attention_params"] = attention_params
                layer_cache["mask_exp"] = mask_exp
                layer_cache["n_tokens"] = n_tokens
                layer_cache["atom_to_token_exp"] = atom_to_token.repeat_interleave(
                    num_diffusion_samples, 0
                )
        else:
            c_base = layer_cache["c_base"]
            attention_params = layer_cache["attention_params"]
            mask_exp = layer_cache["mask_exp"]
            n_tokens = layer_cache["n_tokens"]

        c = c_base

        q = c

        if self.structure_prediction and r_l is not None:
            q = q.repeat_interleave(num_diffusion_samples, 0)
            if pred_r1 is None:
                pred_r1 = torch.zeros_like(r_l)
            r_input = torch.cat([r_l, pred_r1], dim=-1)
            r_to_q = self.coords_linear(r_input)
            q = q + r_to_q

        c = c.repeat_interleave(num_diffusion_samples, 0)

        result = self.atom_transformer(
            q_l=q,
            c_l=c,
            attention_params=attention_params,
            return_intermediates=return_intermediates,
        )
        if return_intermediates:
            q, intermediates = result
        else:
            q = result
            intermediates = []

        q_to_a = F.relu(self.atom_to_token_linear(q))
        if layer_cache is not None and "atom_to_token_exp" in layer_cache:
            atom_to_token_exp = layer_cache["atom_to_token_exp"]
        else:
            atom_to_token_exp = atom_to_token.repeat_interleave(num_diffusion_samples, 0)
        a = scatter_atom_to_token(q_to_a, atom_to_token_exp, n_tokens, atom_mask=mask_exp.bool())

        return a, q, c, attention_params, intermediates


# ===========================================================================
# ESMFold2AtomDecoder
# ===========================================================================


class ESMFold2AtomDecoder(nn.Module):
    """SWA atom decoder with token_to_atom_linear, atom_transformer, norm, output_linear."""

    def __init__(
        self,
        d_atom: int = 128,
        d_token: int = 768,
        n_blocks: int = 3,
        n_heads: int = 4,
        swa_window_size: int = 128,
        expansion_ratio: int = 2,
        spatial_rope_base_frequency: float = 20.0,
        n_spatial_rope_pairs_per_axis: int = 2,
        n_uid_rope_pairs: int = 10,
        uid_rope_base_frequency: float = 10000.0,
    ) -> None:
        super().__init__()
        self.token_to_atom_linear = nn.Linear(d_token, d_atom, bias=False)

        self.atom_transformer = SWAAtomTransformer(
            d_atom=d_atom,
            n_blocks=n_blocks,
            n_heads=n_heads,
            swa_window_size=swa_window_size,
            expansion_ratio=expansion_ratio,
            spatial_rope_base_frequency=spatial_rope_base_frequency,
            n_spatial_rope_pairs_per_axis=n_spatial_rope_pairs_per_axis,
            n_uid_rope_pairs=n_uid_rope_pairs,
            uid_rope_base_frequency=uid_rope_base_frequency,
        )

        self.norm = nn.LayerNorm(d_atom)
        self.output_linear = nn.Linear(d_atom, XYZ_DIMS, bias=False)

    def forward(
        self,
        a_i: Tensor,
        q_l: Tensor,
        c_l: Tensor,
        p_lm: tuple,
        atom_to_token: Tensor,
        atom_attention_mask: Tensor,
        num_diffusion_samples: int = 1,
        return_intermediates: bool = False,
    ) -> tuple[Tensor, list[Tensor]]:
        """Returns (r_update, intermediates)."""
        atom_to_token_exp = atom_to_token.repeat_interleave(num_diffusion_samples, 0)
        a_to_q = self.token_to_atom_linear(a_i)
        a_to_q = gather_token_to_atom(a_to_q, atom_to_token_exp)
        q_l = q_l + a_to_q

        result = self.atom_transformer(
            q_l=q_l,
            c_l=c_l,
            attention_params=p_lm,
            return_intermediates=return_intermediates,
        )
        if return_intermediates:
            q_l, intermediates = result
        else:
            q_l = result
            intermediates = []

        r_l = self.output_linear(self.norm(q_l))
        return r_l, intermediates


# ===========================================================================
# AttentionPairBias (DiffusionTransformer attention block)
# ===========================================================================


class AttentionPairBias(nn.Module):
    """Gated multi-head attention with pair bias conditioning."""

    def __init__(
        self,
        d_model: int,
        d_pair: int,
        num_heads: int,
        d_cond: int | None = None,
        use_conditioning: bool = True,
    ) -> None:
        super().__init__()
        self.d_model = d_model
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads
        self.scale = self.head_dim**-0.5
        d_cond = d_cond or d_model

        if use_conditioning:
            self.adaln = AdaptiveLayerNorm(d_model, d_cond, eps=1e-5)
            self.out_gate = nn.Linear(d_cond, d_model, bias=True)
            # adaln init: weight=0, bias=-2
            nn.init.zeros_(self.out_gate.weight)
            nn.init.constant_(self.out_gate.bias, -2.0)
        else:
            self.pre_norm = nn.LayerNorm(d_model, eps=1e-5)

        self.q_proj = nn.Linear(d_model, d_model, bias=True)
        self.kv_proj = nn.Linear(d_model, 2 * d_model, bias=False)
        self.g_proj = nn.Linear(d_model, d_model, bias=False)
        self.out_proj = nn.Linear(d_model, d_model, bias=False)

        if d_pair > 0:
            self.pair_norm = nn.LayerNorm(d_pair, eps=1e-5)
            self.pair_bias_proj = nn.Linear(d_pair, num_heads, bias=False)

        self._kernel_backend: str | None = None

    def set_kernel_backend(self, backend: str | None) -> None:
        if backend not in _VALID_BACKENDS:
            raise ValueError(f"backend must be one of {_VALID_BACKENDS}, got {backend!r}")
        self._kernel_backend = backend

    def _is_zero_beta(self, beta: Tensor | float) -> bool:
        if isinstance(beta, (int, float)):
            return beta == 0.0
        return bool((beta == 0).all())

    def _can_use_fused_pair_bias(self, z: Tensor, n_queries: int, beta: Tensor | float) -> bool:
        return (
            _fused_active(self, z)
            and z.dim() == 4
            and self._is_zero_beta(beta)
            and hasattr(self, "pair_bias_proj")
            and hasattr(self, "pair_norm")
        )

    def _can_use_cueq_pair_bias(self, z: Tensor, n_queries: int, beta: Tensor | float) -> bool:
        return (
            _cueq_active(self)
            and n_queries > 750
            and z.dim() == 4
            and self._is_zero_beta(beta)
            and hasattr(self, "pair_bias_proj")
        )

    def forward(
        self,
        a: Tensor,
        s: Tensor | None,
        z: Tensor,
        beta: Tensor | float = 0.0,
        attention_mask: Tensor | None = None,
        num_diffusion_samples: int = 1,
    ) -> Tensor:
        bsz, n_queries, d_model = a.shape

        x = self.adaln(a, s) if s is not None else self.pre_norm(a)

        n_keys = x.shape[1]
        q = self.q_proj(x).view(bsz, n_queries, self.num_heads, self.head_dim)
        kv = self.kv_proj(x)
        k, v = kv.chunk(2, dim=-1)
        k = k.view(bsz, n_keys, self.num_heads, self.head_dim)
        v = v.view(bsz, n_keys, self.num_heads, self.head_dim)

        # Expand z for num_diffusion_samples
        if z.dim() == 4 and z.shape[0] != bsz and num_diffusion_samples > 1:
            z = z.repeat_interleave(num_diffusion_samples, dim=0)
        if (
            attention_mask is not None
            and attention_mask.shape[0] != bsz
            and num_diffusion_samples > 1
        ):
            attention_mask = attention_mask.repeat_interleave(num_diffusion_samples, dim=0)

        if self._can_use_fused_pair_bias(z, n_queries, beta):
            kernel_mask = (
                attention_mask
                if attention_mask is not None
                else torch.ones(bsz, n_queries, device=a.device, dtype=torch.bool)
            )
            pair_norm_w = self.pair_norm.weight
            pair_norm_b = (
                self.pair_norm.bias
                if self.pair_norm.bias is not None
                else torch.zeros_like(pair_norm_w)
            )
            z_bf = z if z.dtype == torch.bfloat16 else z.to(torch.bfloat16)
            bias = _fused_pair_bias(  # type: ignore[misc]
                z_bf,
                kernel_mask,
                self.pair_bias_proj.weight,
                num_heads=self.num_heads,
                pair_norm_w=pair_norm_w,
                pair_norm_b=pair_norm_b,
            )  # A has shape (b, h, q, k).
            q_bhqd = q.transpose(1, 2)
            k_bhqd = k.transpose(1, 2)
            v_bhqd = v.transpose(1, 2)
            attn_out = F.scaled_dot_product_attention(
                q_bhqd, k_bhqd, v_bhqd, attn_mask=bias.to(q_bhqd.dtype)
            )
            g = torch.sigmoid(self.g_proj(x)).view(bsz, n_queries, self.num_heads, self.head_dim)
            ctx = g * attn_out.transpose(1, 2)
            out = self.out_proj(ctx.reshape(bsz, n_queries, d_model))
            if s is not None:
                out = torch.sigmoid(self.out_gate(s)) * out
            return out

        if self._can_use_cueq_pair_bias(z, n_queries, beta):
            kernel_mask = (
                attention_mask
                if attention_mask is not None
                else torch.ones(bsz, n_queries, device=a.device, dtype=torch.bool)
            )
            out, _ = _cue_attn_pair_bias(  # type: ignore[misc]
                s=x,
                q=q.transpose(1, 2),
                k=k.transpose(1, 2),
                v=v.transpose(1, 2),
                z=z,
                mask=kernel_mask,
                num_heads=self.num_heads,
                w_proj_z=self.pair_bias_proj.weight,
                w_proj_g=self.g_proj.weight,
                w_proj_o=self.out_proj.weight,
                w_ln_z=self.pair_norm.weight,
                b_ln_z=self.pair_norm.bias,
                return_z_proj=False,
                is_cached_z_proj=False,
            )
        else:
            # Standard attention with pair bias
            g = torch.sigmoid(self.g_proj(x)).view(bsz, n_queries, self.num_heads, self.head_dim)

            logits = torch.einsum("... i h d, ... j h d -> ... i j h", q, k) * self.scale

            pair_bias = self.pair_bias_proj(self.pair_norm(z)) if z.dim() == 4 else z.unsqueeze(-1)
            logits = logits + pair_bias.to(dtype=logits.dtype)

            if attention_mask is not None:
                min_val = torch.finfo(logits.dtype).min
                mask_bias = torch.where(attention_mask.bool()[:, None, :, None], 0.0, min_val)
                logits = logits + mask_bias.to(dtype=logits.dtype)

            attn = torch.softmax(logits, dim=-2).to(dtype=v.dtype)
            ctx = torch.einsum("... i j h, ... j h d -> ... i h d", attn, v)
            ctx = g * ctx
            out = self.out_proj(ctx.reshape(bsz, n_queries, d_model))

        if s is not None:
            out = torch.sigmoid(self.out_gate(s)) * out
        return out


# ===========================================================================
# ConditionedTransitionBlock
# ===========================================================================


class ConditionedTransitionBlock(nn.Module):
    """Conditioned SwiGLU transition with adaptive layer norm."""

    def __init__(
        self,
        d_model: int,
        d_cond: int | None = None,
        transition_multiplier: int = 2,
        use_conditioning: bool = True,
    ) -> None:
        super().__init__()
        d_cond = d_cond or d_model
        hidden = transition_multiplier * d_model

        if use_conditioning:
            self.adaln = AdaptiveLayerNorm(d_model, d_cond, eps=1e-5)
            self.output_gate = nn.Linear(d_cond, d_model, bias=True)
            nn.init.zeros_(self.output_gate.weight)
            nn.init.constant_(self.output_gate.bias, -2.0)
        else:
            self.pre_norm = nn.LayerNorm(d_model, eps=1e-5)

        self.lin_swish = nn.Linear(d_model, 2 * hidden, bias=False)
        self.lin_out = nn.Linear(hidden, d_model, bias=False)

    def forward(self, a: Tensor, s: Tensor | None) -> Tensor:
        x = self.adaln(a, s) if s is not None else self.pre_norm(a)

        swish_a, swish_b = self.lin_swish(x).chunk(2, dim=-1)
        b = F.silu(swish_a) * swish_b
        out = self.lin_out(b)

        if s is not None:
            out = torch.sigmoid(self.output_gate(s)) * out
        return out


# ===========================================================================
# DiffusionTransformer (token transformer)
# ===========================================================================


class DiffusionTransformer(nn.Module):
    """Diffusion denoising transformer with attention pair bias."""

    def __init__(
        self,
        d_model: int,
        d_pair: int,
        num_heads: int,
        num_blocks: int,
        d_cond: int | None = None,
        transition_multiplier: int = 2,
        use_conditioning: bool = True,
    ) -> None:
        super().__init__()
        d_cond = d_cond or d_model

        self.attn_blocks = nn.ModuleList(
            [
                AttentionPairBias(
                    d_model=d_model,
                    d_pair=d_pair,
                    num_heads=num_heads,
                    d_cond=d_cond,
                    use_conditioning=use_conditioning,
                )
                for _ in range(num_blocks)
            ]
        )
        self.transition_blocks = nn.ModuleList(
            [
                ConditionedTransitionBlock(
                    d_model=d_model,
                    d_cond=d_cond,
                    transition_multiplier=transition_multiplier,
                    use_conditioning=use_conditioning,
                )
                for _ in range(num_blocks)
            ]
        )

    def set_kernel_backend(self, backend: str | None) -> None:
        for attn in self.attn_blocks:
            cast(AttentionPairBias, attn).set_kernel_backend(backend)

    def forward(
        self,
        a: Tensor,
        s: Tensor | None,
        z: Tensor,
        beta: Tensor | float = 0.0,
        attention_mask: Tensor | None = None,
        num_diffusion_samples: int = 1,
        return_intermediates: bool = False,
    ) -> tuple[Tensor, list[Tensor]]:
        intermediates: list[Tensor] = []
        x = a
        for attn, transition in zip(self.attn_blocks, self.transition_blocks, strict=True):
            x = x + attn(
                x,
                s,
                z,
                beta,
                attention_mask=attention_mask,
                num_diffusion_samples=num_diffusion_samples,
            )
            x = x + transition(x, s)
            if return_intermediates:
                intermediates.append(x)
        return x, intermediates


# ===========================================================================
# DiffusionConditioning
# ===========================================================================


class DiffusionConditioning(nn.Module):
    """Conditions pair and single representations on noise timestep."""

    def __init__(
        self,
        c_z: int = 256,
        c_s: int = 768,
        c_s_inputs: int = 451,
        sigma_data: float = 16.0,
        fourier_dim: int = 256,
        transition_multiplier: int = 2,
        layer_norm_eps: float = 1e-5,
    ) -> None:
        super().__init__()
        self.sigma_data = float(sigma_data)
        self.c_z = c_z
        self.c_s = c_s
        self.c_s_inputs = c_s_inputs

        self.z_input_norm = nn.LayerNorm(2 * c_z, eps=layer_norm_eps)
        self.z_proj = nn.Linear(2 * c_z, c_z, bias=False)
        self.z_transitions = nn.ModuleList(
            [TransitionLayer(c_z, n=transition_multiplier, eps=layer_norm_eps) for _ in range(2)]
        )

        self.s_input_norm = nn.LayerNorm(c_s_inputs, eps=layer_norm_eps)
        self.s_proj = nn.Linear(c_s_inputs, c_s, bias=False)
        self.fourier = FourierEmbedding(fourier_dim)
        self.noise_norm = nn.LayerNorm(fourier_dim, eps=layer_norm_eps)
        self.noise_proj = nn.Linear(fourier_dim, c_s, bias=False)
        self.s_transitions = nn.ModuleList(
            [TransitionLayer(c_s, n=transition_multiplier, eps=layer_norm_eps) for _ in range(2)]
        )

    def forward(
        self,
        t_hat: Tensor,
        s_inputs: Tensor,
        s_trunk: Tensor | None,
        z_trunk: Tensor,
        relative_position_encoding: Tensor,
        sigma_data: float | None = None,
        num_diffusion_samples: int = 1,
        inference_cache: dict[str, Tensor] | None = None,
    ) -> tuple[Tensor, Tensor]:
        sigma = self.sigma_data if sigma_data is None else float(sigma_data)
        base_batch = z_trunk.shape[0]
        target_batch = base_batch * num_diffusion_samples

        # z conditioning (cached across diffusion steps: independent of t_hat)
        if inference_cache is not None and "z" in inference_cache:
            z = inference_cache["z"]
        else:
            z_rel = relative_position_encoding.to(dtype=torch.float32)
            z = torch.cat([z_trunk.to(dtype=torch.float32), z_rel], dim=-1)
            z = self.z_proj(self.z_input_norm(z))
            with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
                for block in self.z_transitions:
                    z = z + block(z)
            if inference_cache is not None:
                inference_cache["z"] = z

        # s conditioning
        s_inputs_eff = s_inputs
        if s_inputs_eff.shape[0] != target_batch:
            s_inputs_eff = s_inputs_eff.repeat_interleave(num_diffusion_samples, 0)

        s = self.s_proj(self.s_input_norm(s_inputs_eff.to(dtype=torch.float32)))

        # Noise embedding
        t = torch.as_tensor(t_hat, dtype=torch.float32, device=s.device).reshape(-1)
        if t.numel() == 1:
            t = t.expand(target_batch)
        elif t.shape[0] != target_batch:
            t = t.repeat_interleave(num_diffusion_samples, 0)
        t_noise = 0.25 * torch.log((t / sigma).clamp(min=1e-20))
        n = self.fourier(t_noise)
        n = self.noise_proj(self.noise_norm(n))
        s = s + n.unsqueeze(1)

        for block in self.s_transitions:
            s = s + block(s)

        return s, z


# ===========================================================================
# DiffusionModule
# ===========================================================================


class DiffusionModule(nn.Module):
    """Diffusion denoising module for structure prediction."""

    def __init__(
        self,
        c_atom: int = 128,
        c_token: int = 768,
        c_z: int = 256,
        c_s_inputs: int = 451,
        sigma_data: float = 16.0,
        fourier_dim: int = 256,
        atom_num_blocks: int = 3,
        atom_num_heads: int = 4,
        token_num_blocks: int = 12,
        token_num_heads: int = 16,
        transition_multiplier: int = 2,
        swa_window_size: int = 128,
        spatial_rope_base_frequency: float = 20.0,
        n_spatial_rope_pairs_per_axis: int = 2,
        n_uid_rope_pairs: int = 10,
        uid_rope_base_frequency: float = 10000.0,
    ) -> None:
        super().__init__()
        self.sigma_data = float(sigma_data)

        self.conditioning = DiffusionConditioning(
            c_z=c_z,
            c_s=c_token,  # conditioning s output is c_token
            c_s_inputs=c_s_inputs,
            sigma_data=sigma_data,
            fourier_dim=fourier_dim,
            transition_multiplier=transition_multiplier,
        )

        # Atom encoder (structure_prediction=True, with coords_linear)
        self.atom_encoder = ESMFold2AtomEncoder(
            d_atom=c_atom,
            d_token=c_token,
            n_blocks=atom_num_blocks,
            n_heads=atom_num_heads,
            swa_window_size=swa_window_size,
            expansion_ratio=2,
            structure_prediction=True,
            spatial_rope_base_frequency=spatial_rope_base_frequency,
            n_spatial_rope_pairs_per_axis=n_spatial_rope_pairs_per_axis,
            n_uid_rope_pairs=n_uid_rope_pairs,
            uid_rope_base_frequency=uid_rope_base_frequency,
        )

        # Atom decoder
        self.atom_decoder = ESMFold2AtomDecoder(
            d_atom=c_atom,
            d_token=c_token,
            n_blocks=atom_num_blocks,
            n_heads=atom_num_heads,
            swa_window_size=swa_window_size,
            expansion_ratio=2,
            spatial_rope_base_frequency=spatial_rope_base_frequency,
            n_spatial_rope_pairs_per_axis=n_spatial_rope_pairs_per_axis,
            n_uid_rope_pairs=n_uid_rope_pairs,
            uid_rope_base_frequency=uid_rope_base_frequency,
        )

        self.s_to_token = nn.Linear(c_token, c_token, bias=False)
        nn.init.zeros_(self.s_to_token.weight)

        # Token transformer (DiffusionTransformer with pair bias)
        self.token_transformer = DiffusionTransformer(
            d_model=c_token,
            d_pair=c_z,
            num_heads=token_num_heads,
            num_blocks=token_num_blocks,
            d_cond=c_token,
            transition_multiplier=transition_multiplier,
            use_conditioning=True,
        )

        self.s_step_norm = nn.LayerNorm(c_token)
        self.token_norm = nn.LayerNorm(c_token)

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

    def forward(
        self,
        x_noisy: Tensor,
        t_hat: Tensor,
        ref_pos: Tensor,
        ref_charge: Tensor,
        ref_mask: Tensor,
        ref_element: Tensor,
        ref_atom_name_chars: Tensor,
        ref_space_uid: Tensor,
        tok_idx: Tensor,
        s_inputs: Tensor,
        s_trunk: Tensor | None,
        z_trunk: Tensor,
        relative_position_encoding: Tensor,
        asym_id: Tensor,
        residue_index: Tensor,
        entity_id: Tensor,
        token_index: Tensor,
        sym_id: Tensor,
        sigma_data: float | None = None,
        token_attention_mask: Tensor | None = None,
        num_diffusion_samples: int = 1,
        return_token_repr: bool = False,
        return_atom_repr: bool = False,
        inference_cache: dict[str, Tensor] | None = None,
    ) -> dict[str, Tensor | None]:
        bsz = x_noisy.shape[0]
        sigma = self.sigma_data if sigma_data is None else float(sigma_data)
        t = torch.as_tensor(t_hat, dtype=torch.float32, device=x_noisy.device).reshape(-1)
        if t.numel() == 1:
            t = t.expand(bsz)

        # Step 1: conditioning (pair z is cached across diffusion steps)
        s, z = self.conditioning(
            t_hat=t,
            s_inputs=s_inputs,
            s_trunk=s_trunk,
            z_trunk=z_trunk,
            relative_position_encoding=relative_position_encoding,
            sigma_data=sigma,
            num_diffusion_samples=num_diffusion_samples,
            inference_cache=inference_cache,
        )

        # Step 2: normalize noisy coords
        denom = torch.sqrt(t * t + sigma * sigma)
        r_noisy = x_noisy / denom[:, None, None]

        # Step 3: atom encoder
        a, q_skip, c_skip, p_skip, enc_intermediates = self.atom_encoder(
            ref_pos=ref_pos,
            atom_attention_mask=ref_mask,
            ref_space_uid=ref_space_uid,
            ref_charge=ref_charge,
            ref_element=ref_element,
            ref_atom_name_chars=ref_atom_name_chars,
            atom_to_token=tok_idx,
            r_l=r_noisy,
            s_i=s_trunk,
            num_diffusion_samples=num_diffusion_samples,
            return_intermediates=return_atom_repr,
            inference_cache=inference_cache,
        )

        # Step 4: add conditioned s
        a = a + self.s_to_token(self.s_step_norm(s))

        # Step 5: token transformer
        a, _ = self.token_transformer(
            a,
            s,
            z,
            beta=0.0,
            attention_mask=token_attention_mask,
            num_diffusion_samples=num_diffusion_samples,
        )

        # Step 6: token norm
        a = self.token_norm(a)

        # Step 7: atom decoder
        r_update, dec_intermediates = self.atom_decoder(
            a_i=a,
            q_l=q_skip,
            c_l=c_skip,
            p_lm=p_skip,
            atom_to_token=tok_idx,
            atom_attention_mask=ref_mask,
            num_diffusion_samples=num_diffusion_samples,
            return_intermediates=return_atom_repr,
        )

        # Step 8: compute denoised output
        sigma2 = sigma * sigma
        t2 = t * t
        out = (sigma2 / (sigma2 + t2))[:, None, None] * x_noisy
        out = out + ((sigma * t) / torch.sqrt(sigma2 + t2))[:, None, None] * r_update

        # Collect atom intermediates from encoder + decoder
        atom_intermediates: Tensor | None = None
        if return_atom_repr:
            all_ints = enc_intermediates + dec_intermediates
            if all_ints:
                atom_intermediates = torch.stack(all_ints, dim=2)

        return {
            "x_denoised": out,
            "token_repr": a if return_token_repr else None,
            "atom_intermediates": atom_intermediates,
        }


# ===========================================================================
# DiffusionStructureHead
# ===========================================================================


class DiffusionStructureHead(nn.Module):
    """Wrapper around DiffusionModule with diffusion sampling."""

    def __init__(self, config: ESMFold2Config) -> None:
        super().__init__()
        dm = config.structure_head.diffusion_module
        swa_cfg = config.inputs.atom_encoder
        sh = config.structure_head

        self.diffusion_module = DiffusionModule(
            c_atom=dm.c_atom,
            c_token=dm.c_token,
            c_z=dm.c_z,
            c_s_inputs=dm.c_s_inputs,
            sigma_data=dm.sigma_data,
            fourier_dim=dm.fourier_dim,
            atom_num_blocks=dm.atom_num_blocks,
            atom_num_heads=dm.atom_num_heads,
            token_num_blocks=dm.token_num_blocks,
            token_num_heads=dm.token_num_heads,
            transition_multiplier=dm.transition_multiplier,
            swa_window_size=swa_cfg.swa_window_size,
            spatial_rope_base_frequency=swa_cfg.spatial_rope_base_frequency,
            n_spatial_rope_pairs_per_axis=swa_cfg.n_spatial_rope_pairs_per_axis,
            n_uid_rope_pairs=swa_cfg.n_uid_rope_pairs,
            uid_rope_base_frequency=swa_cfg.uid_rope_base_frequency,
        )

        # Sampling hyperparameters
        self.sigma_data = dm.sigma_data
        self.gamma_0 = sh.gamma_0
        self.gamma_min = sh.gamma_min
        self.noise_scale = sh.noise_scale
        self.step_scale = sh.step_scale
        self.inference_s_max = sh.inference_s_max
        self.inference_s_min = sh.inference_s_min
        self.inference_p = sh.inference_p
        self.inference_num_steps = sh.inference_num_steps

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

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    def inference_noise_schedule(
        self, num_steps: int | None = None, device: torch.device | None = None
    ) -> Tensor:
        """Karras power-law noise schedule."""
        steps = self.inference_num_steps if num_steps is None else int(num_steps)
        if steps == 1:
            return torch.tensor(
                [self.inference_s_max * self.sigma_data, 0.0],
                device=device,
                dtype=torch.float32,
            )
        p = float(self.inference_p)
        inv_p = 1.0 / p
        k = torch.arange(steps, device=device, dtype=torch.float32)
        base = self.inference_s_max**inv_p + (k / (steps - 1)) * (
            self.inference_s_min**inv_p - self.inference_s_max**inv_p
        )
        schedule = self.sigma_data * base.pow(p)
        return F.pad(schedule, (0, 1), value=0.0)

    @staticmethod
    def _random_rotations(n: int, dtype: torch.dtype, device: torch.device) -> Tensor:
        q = torch.randn((n, 4), dtype=dtype, device=device)
        scale = torch.sqrt((q * q).sum(dim=1))
        signs = torch.where(q[:, 0] < 0, -scale, scale)
        q = q / signs[:, None]
        r, i, j, k = torch.unbind(q, dim=-1)
        two_s = 2.0 / (q * q).sum(dim=-1)
        return torch.stack(
            (
                1 - two_s * (j * j + k * k),
                two_s * (i * j - k * r),
                two_s * (i * k + j * r),
                two_s * (i * j + k * r),
                1 - two_s * (i * i + k * k),
                two_s * (j * k - i * r),
                two_s * (i * k - j * r),
                two_s * (j * k + i * r),
                1 - two_s * (i * i + j * j),
            ),
            dim=-1,
        ).reshape(n, 3, 3)

    def _center_random_augmentation(
        self, x: Tensor, atom_mask: Tensor, second_coords: Tensor | None = None
    ) -> tuple[Tensor, Tensor | None]:
        """Algorithm 19: center + random rotation + translation."""
        bsz = x.shape[0]
        mask = atom_mask.unsqueeze(-1)  # M has shape (b, a, 1).
        denom = mask.sum(dim=1, keepdim=True).clamp(min=1)
        mean = (x * mask).sum(dim=1, keepdim=True) / denom
        x = x - mean
        if second_coords is not None:
            second_coords = second_coords - mean

        r = self._random_rotations(bsz, x.dtype, x.device)
        x = torch.einsum("bmd,bds->bms", x, r)
        if second_coords is not None:
            second_coords = torch.einsum("bmd,bds->bms", second_coords, r)

        t = torch.randn_like(x[:, 0:1, :])
        x = x + t
        if second_coords is not None:
            second_coords = second_coords + t
        return x, second_coords

    @staticmethod
    def _weighted_rigid_align(x: Tensor, x_gt: Tensor, w: Tensor, mask: Tensor) -> Tensor:
        """Kabsch alignment: align x to x_gt with weights w."""
        w = (mask * w).unsqueeze(-1)  # W has shape (b, n, 1).
        denom = w.sum(dim=-2, keepdim=True).clamp(min=1e-8)
        mu = (x * w).sum(dim=-2, keepdim=True) / denom
        mu_gt = (x_gt * w).sum(dim=-2, keepdim=True) / denom
        x_c = x - mu
        xgt_c = x_gt - mu_gt
        covariance = torch.einsum("bni,bnj->bij", w * xgt_c, x_c)
        covariance_f32 = covariance.float()
        u, _, vh = torch.linalg.svd(
            covariance_f32, driver="gesvd" if covariance_f32.is_cuda else None
        )
        det = torch.linalg.det(u @ vh)
        ones = torch.ones_like(det)
        rotation = (u @ torch.diag_embed(torch.stack([ones, ones, det], dim=-1)) @ vh).to(
            covariance.dtype
        )
        return x_c @ rotation.transpose(-1, -2) + mu_gt

    # ------------------------------------------------------------------
    # Sampling
    # ------------------------------------------------------------------

    @torch.inference_mode()
    def sample(
        self,
        z_trunk: Tensor,
        s_inputs: Tensor,
        s_trunk: Tensor | None,
        relative_position_encoding: Tensor,
        ref_pos: Tensor,
        ref_charge: Tensor,
        ref_mask: Tensor,
        ref_element: Tensor,
        ref_atom_name_chars: Tensor,
        ref_space_uid: Tensor,
        tok_idx: Tensor,
        asym_id: Tensor,
        residue_index: Tensor,
        entity_id: Tensor,
        token_index: Tensor,
        sym_id: Tensor,
        token_attention_mask: Tensor | None = None,
        num_diffusion_samples: int = 1,
        num_sampling_steps: int | None = None,
        max_inference_sigma: float | None = 256.0,
        noise_scale: float | None = None,
        step_scale: float | None = None,
        return_atom_repr: bool = False,
        use_inference_cache: bool = True,
        denoising_early_exit_rmsd: float | None = None,
    ) -> dict[str, Tensor | None]:
        """Diffusion sampling (Algorithm 18).

        ``num_sampling_steps`` is the number of denoising steps actually run.
        When ``max_inference_sigma`` is set, the Karras schedule built with
        ``num_sampling_steps`` entries would lose its high-sigma tail to the cap,
        so we inflate the underlying schedule length here to land back at the
        requested step count post-truncation.
        """
        n_atoms = tok_idx.shape[1]
        device = s_inputs.device
        target_batch = s_inputs.shape[0] * num_diffusion_samples

        inference_cache: dict[str, Tensor] | None = {} if use_inference_cache else None

        steps = self.inference_num_steps if num_sampling_steps is None else int(num_sampling_steps)

        schedule = self.inference_noise_schedule(steps, device)
        if max_inference_sigma is not None:
            schedule = schedule[schedule <= float(max_inference_sigma)]
            schedule = F.pad(schedule, (1, 0), value=float(max_inference_sigma))

        lam = self.noise_scale if noise_scale is None else float(noise_scale)
        eta = self.step_scale if step_scale is None else float(step_scale)

        x = schedule[0] * torch.randn(target_batch, n_atoms, 3, device=device, dtype=torch.float32)
        atom_mask = ref_mask.repeat_interleave(num_diffusion_samples, 0).float()

        gammas = torch.where(
            schedule > self.gamma_min,
            torch.full_like(schedule, self.gamma_0),
            torch.zeros_like(schedule),
        )

        x_denoised_prev: Tensor | None = None
        token_repr: Tensor | None = None
        diff_atom_intermediates: Tensor | None = None

        step_pairs = list(zip(schedule[:-1], schedule[1:], gammas[1:], strict=True))
        num_steps = len(step_pairs)

        for step_idx, (sigma_tm, sigma_t, gamma) in enumerate(step_pairs):
            x, x_denoised_prev = self._center_random_augmentation(
                x, atom_mask, second_coords=x_denoised_prev
            )

            sigma_tm_val = float(sigma_tm.item())
            t_hat_val = sigma_tm_val * (1.0 + float(gamma.item()))
            eps_std = lam * max(t_hat_val**2 - sigma_tm_val**2, 0.0) ** 0.5
            x_noisy = x + eps_std * torch.randn_like(x)

            is_last_step = step_idx == num_steps - 1
            request_atom_repr = return_atom_repr and (
                is_last_step or denoising_early_exit_rmsd is not None
            )

            dm_out = self.diffusion_module(
                x_noisy=x_noisy,
                t_hat=torch.full((target_batch,), t_hat_val, device=device, dtype=torch.float32),
                ref_pos=ref_pos,
                ref_charge=ref_charge,
                ref_mask=ref_mask,
                ref_element=ref_element,
                ref_atom_name_chars=ref_atom_name_chars,
                ref_space_uid=ref_space_uid,
                tok_idx=tok_idx,
                s_inputs=s_inputs,
                s_trunk=s_trunk,
                z_trunk=z_trunk,
                relative_position_encoding=relative_position_encoding,
                asym_id=asym_id,
                residue_index=residue_index,
                entity_id=entity_id,
                token_index=token_index,
                sym_id=sym_id,
                token_attention_mask=token_attention_mask,
                num_diffusion_samples=num_diffusion_samples,
                return_token_repr=True,
                return_atom_repr=request_atom_repr,
                inference_cache=inference_cache,
            )

            x_denoised = dm_out["x_denoised"]
            token_repr = dm_out["token_repr"]
            if request_atom_repr:
                diff_atom_intermediates = dm_out.get("atom_intermediates")

            # Reverse diffusion alignment (Kabsch)
            with torch.autocast(device_type="cuda", enabled=False):
                x_noisy = self._weighted_rigid_align(
                    x_noisy.float(), x_denoised.float(), atom_mask, atom_mask
                )
            x_noisy = x_noisy.to(dtype=x_denoised.dtype)

            # ODE/SDE step
            sigma_t_val = float(sigma_t.item())
            denoised_over_sigma = (x_noisy - x_denoised) / t_hat_val
            x = x_noisy + eta * (sigma_t_val - t_hat_val) * denoised_over_sigma

            # Denoising early-exit: stop when consecutive predictions converge
            if (
                denoising_early_exit_rmsd is not None
                and x_denoised_prev is not None
                and step_idx >= 1
            ):
                with torch.autocast(device_type="cuda", enabled=False):
                    aligned = self._weighted_rigid_align(
                        x_denoised_prev.float(),
                        x_denoised.float(),
                        atom_mask,
                        atom_mask,
                    )
                diff = (x_denoised.float() - aligned) * atom_mask.unsqueeze(-1)
                per_sample_rmsd = (
                    diff.pow(2).sum(dim=(-1, -2)) / atom_mask.sum(dim=-1).clamp(min=1)
                ).sqrt()
                if per_sample_rmsd.max().item() < denoising_early_exit_rmsd:
                    x = x_denoised
                    x_denoised_prev = x_denoised
                    break

            x_denoised_prev = x_denoised

        result: dict[str, Tensor | None] = {
            "sample_atom_coords": x,
            "diff_token_repr": token_repr,
        }
        if return_atom_repr:
            result["diff_atom_intermediates"] = diff_atom_intermediates
        return result