rc-foundry 0.1.1__py3-none-any.whl

rf3/model/layers/attention.py

@@ -0,0 +1,313 @@
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from jaxtyping import Float
+from opt_einsum import contract as einsum
+from rf3.util_module import init_lecun_normal
+
+from foundry import SHOULD_USE_CUEQUIVARIANCE
+from foundry.training.checkpoint import activation_checkpointing
+
+if SHOULD_USE_CUEQUIVARIANCE:
+    import cuequivariance_torch as cuet
+
+
+class TriangleAttention(nn.Module):
+    """Implementation of Triangle Attention from AlphaFold3.
+
+    Routes to either cuEquivariance or vanilla implementation based on configuration.
+
+    Args:
+        pair: Pair representation tensor of shape (B, L, L, d_pair)
+
+    Returns:
+        Updated pair representation tensor of shape (B, L, L, d_pair)
+
+    """
+
+    def __init__(
+        self,
+        d_pair,
+        n_head=4,
+        d_hidden=32,
+        p_drop=0.1,  # noqa: E402
+        start_node=True,
+        use_cuequivariance=True,
+    ):
+        super(TriangleAttention, self).__init__()
+
+        self.norm = nn.LayerNorm(d_pair)
+
+        self.to_q = nn.Linear(d_pair, n_head * d_hidden, bias=False)
+        self.to_k = nn.Linear(d_pair, n_head * d_hidden, bias=False)
+        self.to_v = nn.Linear(d_pair, n_head * d_hidden, bias=False)
+        self.to_b = nn.Linear(d_pair, n_head, bias=False)
+        self.to_g = nn.Linear(d_pair, n_head * d_hidden)
+        self.to_out = nn.Linear(n_head * d_hidden, d_pair)
+
+        self.scaling = 1 / math.sqrt(d_hidden)
+
+        self.h = n_head
+        self.dim = d_hidden
+        self.start_node = start_node
+
+        self.use_cuequivariance = use_cuequivariance
+
+        self.reset_parameter()
+
+    def reset_parameter(self):
+        # query/key/value projection: Glorot uniform / Xavier uniform
+        nn.init.xavier_uniform_(self.to_q.weight)
+        nn.init.xavier_uniform_(self.to_k.weight)
+        nn.init.xavier_uniform_(self.to_v.weight)
+
+        # bias: normal distribution
+        self.to_b = init_lecun_normal(self.to_b)
+
+        # gating: zero weights, one biases (mostly open gate at the begining)
+        nn.init.zeros_(self.to_g.weight)
+        nn.init.ones_(self.to_g.bias)
+
+        # to_out: right before residual connection: zero initialize -- to make it sure residual operation is same to the Identity at the begining
+        nn.init.zeros_(self.to_out.weight)
+        nn.init.zeros_(self.to_out.bias)
+
+    @activation_checkpointing
+    def forward(self, pair):
+        """Forward pass of triangle attention."""
+        pair = self.norm(pair)
+        bias = self.to_b(pair)  # (B, L, L, h)
+
+        if not self.start_node:
+            pair = rearrange(pair, "b i j d -> b j i d")
+
+        # Route to appropriate implementation
+        if self.use_cuequivariance and SHOULD_USE_CUEQUIVARIANCE:
+            out = self._forward_cuequivariance(pair, bias)
+        else:
+            out = self._forward_vanilla(pair, bias)
+
+        if not self.start_node:
+            out = rearrange(out, "b i j d -> b j i d")
+
+        # output projection
+        out = self.to_out(out)
+        return out
+
+    def _forward_cuequivariance(self, pair, bias):
+        """cuEquivariance triangle attention implementation."""
+        # Handle autocast conversion
+        if torch.is_autocast_enabled():
+            dtype = torch.get_autocast_dtype("cuda")
+            pair = pair.to(dtype=dtype)
+            bias = bias.to(dtype=dtype)
+
+        assert (
+            pair.dtype == torch.bfloat16 and bias.dtype == torch.bfloat16
+        ), f"cuEquivariance requires bfloat16 inputs (got pair={pair.dtype}, bias={bias.dtype})"
+
+        # Gate computation
+        gate = torch.sigmoid(self.to_g(pair))  # (B, L, L, h*dim)
+
+        # Project and reshape to cuEquivariance format: (B, L, H, L, D)
+        query = rearrange(self.to_q(pair), "b i j (h d) -> b i h j d", h=self.h)
+        key = rearrange(self.to_k(pair), "b i k (h d) -> b i h k d", h=self.h)
+        value = rearrange(self.to_v(pair), "b i k (h d) -> b i h k d", h=self.h)
+
+        # Bias: (B, L, L, H) -> (B, 1, H, L, L)
+        bias_cueq = rearrange(bias, "b i j h -> b 1 h i j")
+
+        # Call cuEquivariance triangle attention
+        out_cueq = cuet.triangle_attention(
+            query, key, value, bias=bias_cueq, scale=self.scaling
+        )
+
+        # Reshape back: (B, L, H, L, D) -> (B, L, L, H*D)
+        out = rearrange(out_cueq, "b i h j d -> b i j (h d)")
+        out = gate * out  # gated attention
+        return out
+
+    def _forward_vanilla(self, pair, bias):
+        """Vanilla PyTorch triangle attention implementation."""
+        B, L = pair.shape[:2]
+
+        # Gate computation
+        gate = torch.sigmoid(self.to_g(pair))  # (B, L, L, h*dim)
+
+        # Project and reshape to vanilla format: (B, L, L, H, D)
+        query = self.to_q(pair).reshape(B, L, L, self.h, -1)
+        key = self.to_k(pair).reshape(B, L, L, self.h, -1)
+        value = self.to_v(pair).reshape(B, L, L, self.h, -1)
+
+        query = query * self.scaling
+
+        attn = einsum("bijhd,bikhd->bijkh", query, key)
+        attn = attn + bias.unsqueeze(1).expand(-1, L, -1, -1, -1)  # (bijkh)
+        attn = F.softmax(attn, dim=-2)
+
+        out = einsum("bijkh,bikhd->bijhd", attn, value).reshape(B, L, L, -1)
+        out = gate * out  # gated attention
+
+        return out
+
+
+class TriangleMultiplication(nn.Module):
+    """Implementation of Triangle Multiplicative Update from AlphaFold3.
+
+    Routes to either cuEquivariance or naive implementation based on configuration.
+
+    Args:
+        d_pair: Pair representation dimension (must equal d_hidden for cuEquivariance)
+        d_hidden: Hidden dimension (must equal d_pair for cuEquivariance)
+        direction: "outgoing" or "incoming" triangle multiplication direction
+        bias: Whether to use bias in normalization layers
+        use_cuequivariance: Whether to use cuEquivariance fused kernel when available
+    """
+
+    def __init__(
+        self,
+        d_pair,
+        d_hidden=None,
+        direction="outgoing",
+        bias=True,
+        use_cuequivariance=True,
+    ):
+        super(TriangleMultiplication, self).__init__()
+
+        # Set d_hidden to d_pair if not specified
+        if d_hidden is None:
+            d_hidden = d_pair
+
+        self.d_pair = d_pair
+        self.d_hidden = d_hidden
+
+        # Validate direction parameter
+        if direction not in ["outgoing", "incoming"]:
+            raise ValueError(
+                f"direction must be 'outgoing' or 'incoming', got '{direction}'"
+            )
+        self.direction = direction
+
+        self.use_cuequivariance = use_cuequivariance
+
+        if self.use_cuequivariance:
+            # cuEquivariance kernel requires d_pair == d_hidden...
+            assert (
+                d_pair == d_hidden
+            ), "cuEquivariance triangle multiplication requires d_pair == d_hidden"
+            # ... and d_pair must be a multiple of 32
+            assert (
+                d_pair % 32 == 0
+            ), "cuEquivariance triangle multiplication requires d_pair to be a multiple of 32"
+
+        # Input normalization (optional bias)
+        self.norm_in = nn.LayerNorm(d_pair, bias=bias)
+
+        # Input projections: combine left and right projections (2*d_hidden, d_pair) (no bias)
+        self.p_in = nn.Linear(d_pair, 2 * d_hidden, bias=False)
+
+        # Input gating: combine left and right gates (2*d_hidden, d_pair) (no bias)
+        self.g_in = nn.Linear(d_pair, 2 * d_hidden, bias=False)
+
+        # Output normalization (optional bias)
+        self.norm_out = nn.LayerNorm(d_hidden, bias=bias)
+
+        # Output projection (no bias)
+        self.p_out = nn.Linear(d_hidden, d_pair, bias=False)
+
+        # Output gating (no bias)
+        self.g_out = nn.Linear(d_pair, d_pair, bias=False)
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        # Input projections: lecun normal distribution for regular linear weights
+        self.p_in = init_lecun_normal(self.p_in)
+
+        # We use default PyTorch initialization for the other parameters, as in AF-3 they do not specify their
+        # weight initialization schemes. Without bias, e.g., the gate initialization from AF-2 is not correct.
+
+    def forward(
+        self, pair: Float[torch.Tensor, "B N N D"]
+    ) -> Float[torch.Tensor, "B N N D"]:
+        """Forward pass of triangle multiplication."""
+        # Route to appropriate implementation
+        if self.use_cuequivariance and SHOULD_USE_CUEQUIVARIANCE:
+            return self._forward_cuequivariance(pair)
+        else:
+            return self._forward_vanilla(pair)
+
+    def _forward_vanilla(
+        self, pair: Float[torch.Tensor, "B N N D"]
+    ) -> Float[torch.Tensor, "B N N D"]:
+        """Vanilla PyTorch triangle multiplication implementation."""
+        B, L = pair.shape[:2]
+
+        # Input normalization
+        pair_norm = self.norm_in(pair)
+
+        # Input projections: get combined output and split
+        p_combined = self.p_in(pair_norm)  # (B, L, L, 2*d_hidden)
+        left = p_combined[..., : self.d_hidden]  # (B, L, L, d_hidden)
+        right = p_combined[..., self.d_hidden :]  # (B, L, L, d_hidden)
+
+        # Input gating: get combined output and split
+        g_combined = self.g_in(pair_norm)  # (B, L, L, 2*d_hidden)
+        left_gate = torch.sigmoid(g_combined[..., : self.d_hidden])
+        right_gate = torch.sigmoid(g_combined[..., self.d_hidden :])
+
+        # Apply gating
+        left = left_gate * left
+        right = right_gate * right
+
+        # Triangle multiplication based on direction
+        if self.direction == "outgoing":
+            out = torch.einsum("bikd,bjkd->bijd", left, right / float(L))
+        else:  # incoming
+            out = torch.einsum("bkid,bkjd->bijd", left, right / float(L))
+
+        # Output normalization
+        out = self.norm_out(out)
+
+        # Output projection
+        out = self.p_out(out)
+
+        # Output gating
+        gate = torch.sigmoid(self.g_out(pair_norm))
+        out = gate * out
+
+        return out
+
+    def _forward_cuequivariance(
+        self, pair: Float[torch.Tensor, "B N N D"]
+    ) -> Float[torch.Tensor, "B N N D"]:
+        """cuEquivariance triangle multiplication implementation."""
+        # Handle autocast conversion
+        # (Use bfloat16 for optimal performance)
+        if torch.is_autocast_enabled():
+            dtype = torch.get_autocast_dtype("cuda")
+            pair = pair.to(dtype=dtype)
+
+        assert (
+            pair.dtype == torch.bfloat16
+        ), "cuEquivariance requires bfloat16 inputs for optimal performance"
+
+        output = cuet.triangle_multiplicative_update(
+            x=pair,
+            direction=self.direction,
+            mask=None,
+            norm_in_weight=self.norm_in.weight,
+            norm_in_bias=self.norm_in.bias,
+            p_in_weight=self.p_in.weight,  # (2*d_hidden, d_pair)
+            g_in_weight=self.g_in.weight,  # (2*d_hidden, d_pair)
+            norm_out_weight=self.norm_out.weight,
+            norm_out_bias=self.norm_out.bias,
+            p_out_weight=self.p_out.weight,  # (d_pair, d_pair) since d_hidden == d_pair
+            g_out_weight=self.g_out.weight,  # (d_pair, d_pair)
+            eps=1e-5,
+        )
+
+        return output

rf3/model/layers/layer_utils.py

@@ -0,0 +1,127 @@
+from functools import partial
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.nn.functional import silu
+
+from foundry.training.checkpoint import activation_checkpointing
+
+linearNoBias = partial(torch.nn.Linear, bias=False)
+
+
+def collapse(x, L):
+    return x.reshape((L, x.numel() // L))
+
+
+class MultiDimLinear(nn.Linear):
+    def __init__(self, in_features, out_shape, **kwargs):
+        self.out_shape = out_shape
+        out_features = np.prod(out_shape)
+        super().__init__(in_features, out_features, **kwargs)
+        self.reset_parameters()
+
+    def reset_parameters(self, **kwargs) -> None:
+        super().reset_parameters()
+        nn.init.xavier_uniform_(self.weight)
+
+    def forward(self, x):
+        out = super().forward(x)
+        return out.reshape(x.shape[:-1] + self.out_shape)
+
+
+class LinearBiasInit(nn.Linear):
+    def __init__(self, *args, biasinit, **kwargs):
+        assert biasinit == -2.0  # Sanity check
+        self.biasinit = biasinit
+        super().__init__(*args, **kwargs)
+
+    def reset_parameters(self) -> None:
+        super().reset_parameters()
+        self.bias.data.fill_(self.biasinit)
+
+
+class Transition(nn.Module):
+    def __init__(self, n, c):
+        super().__init__()
+        self.layer_norm_1 = nn.LayerNorm(c)
+        self.linear_1 = linearNoBias(c, n * c)
+        self.linear_2 = linearNoBias(c, n * c)
+        self.linear_3 = linearNoBias(n * c, c)
+
+    @activation_checkpointing
+    def forward(
+        self,
+        X,
+    ):
+        X = self.layer_norm_1(X)
+        A = self.linear_1(X)
+        B = self.linear_2(X)
+        X = self.linear_3(silu(A) * B)
+        return X
+
+
+class AdaLN(nn.Module):
+    def __init__(self, c_a, c_s, n=2):
+        super().__init__()
+        self.ln_a = nn.LayerNorm(normalized_shape=(c_a,), elementwise_affine=False)
+        self.ln_s = nn.LayerNorm(normalized_shape=(c_s,), bias=False)
+        self.to_gain = nn.Sequential(
+            nn.Linear(c_s, c_a),
+            nn.Sigmoid(),
+        )
+        self.to_bias = linearNoBias(c_s, c_a)
+
+    def forward(
+        self,
+        Ai,  # [B, I, C_a]
+        Si,  # [B, I, C_s]
+    ):
+        """
+        Output:
+            [B, I, C_a]
+        """
+        Ai = self.ln_a(Ai)
+        Si = self.ln_s(Si)
+        return self.to_gain(Si) * Ai + self.to_bias(Si)
+
+
+def create_batch_dimension_if_not_present(batched_n_dim):
+    """
+    Decorator for adapting a function which expects batched arguments with ndim `batched_n_dim` also
+    accept unbatched arguments.
+    """
+
+    def wrap(f):
+        def _wrap(arg):
+            inserted_batch_dim = False
+            if arg.ndim == batched_n_dim - 1:
+                arg = arg[None]
+                inserted_batch_dim = True
+            elif arg.ndim == batched_n_dim:
+                pass
+            else:
+                raise Exception(
+                    f"arg must have {batched_n_dim - 1} or {batched_n_dim} dimensions, got shape {arg.shape=}"
+                )
+            o = f(arg)
+
+            if inserted_batch_dim:
+                assert o.shape[0] == 1, f"{o.shape=}[0] != 1"
+                return o[0]
+            return o
+
+        return _wrap
+
+    return wrap
+
+
+def unpack_args_for_checkpointing(arg_names):
+    def wrap(f):
+        def _wrap(*args):
+            f = args[0]
+            return f(**dict(zip(arg_names, args)))
+
+        return _wrap
+
+    return wrap

rf3/model/layers/mlff.py

@@ -0,0 +1,118 @@
+import torch
+import torch.nn as nn
+from einops import rearrange
+from jaxtyping import Float
+
+
+class ConformerEmbeddingWeightedAverage(nn.Module):
+    """Learned weighted average of reference conformer embeddings.
+
+    Args:
+        atom_level_embedding_dim: Dimension of the input atom-level embeddings (default: 384, for EGRET)
+        c_atompair: Dimension of the atom-pair embeddings (default: 16)
+        c_atom: Dimension of the output atom embeddings (default: 128)
+        n_conformers: Number of conformers to expect (default: 8)
+        dropout_rate: Dropout rate for regularization (default: 0.1)
+        use_layer_norm: Whether to apply layer normalization to the output (default: True)
+    """
+
+    def __init__(
+        self,
+        atom_level_embedding_dim: int,
+        c_atompair: int,
+        c_atom: int,
+        n_conformers: int = 8,
+        dropout_rate: float = 0.1,
+        use_layer_norm: bool = True,
+    ):
+        super().__init__()
+
+        self.n_conformers = n_conformers
+        self.atom_level_embedding_dim = atom_level_embedding_dim
+
+        # Downcast MLP from atom_level_embedding_dim to c_atompair
+        self.process_atom_level_embedding = nn.Sequential(
+            nn.Linear(atom_level_embedding_dim, atom_level_embedding_dim // 2),
+            nn.ReLU(),
+            nn.Dropout(dropout_rate),
+            nn.Linear(atom_level_embedding_dim // 2, atom_level_embedding_dim // 4),
+            nn.ReLU(),
+            nn.Dropout(dropout_rate),
+            nn.Linear(atom_level_embedding_dim // 4, atom_level_embedding_dim // 8),
+            nn.ReLU(),
+            nn.Dropout(dropout_rate),
+            nn.Linear(atom_level_embedding_dim // 8, c_atompair),
+        )
+
+        # Final MLP to convert from (n_conformers * c_atompair) to c_atom
+        self.conformers_to_atom_single_embedding = nn.Sequential(
+            nn.Linear(n_conformers * c_atompair, c_atom, bias=False),
+            nn.LayerNorm(c_atom) if use_layer_norm else nn.Identity(),
+        )
+
+        # Zero-init the final linear layer to ensure the model starts with identity function (output ≈ 0)
+        nn.init.zeros_(self.conformers_to_atom_single_embedding[0].weight)
+
+    def forward(
+        self,
+        atom_level_embeddings: Float[
+            torch.Tensor, "n_conformers n_atom atom_level_embedding_dim"
+        ],
+    ) -> Float[torch.Tensor, "n_atom c_atom"]:
+        """Forward pass: process atom-level embeddings and return the processed result.
+
+        Args:
+            atom_level_embeddings: Input tensor of shape [n_conformers, n_atom, atom_level_embedding_dim]
+
+        Returns:
+            Processed tensor of shape [n_atom, c_atom] ready for residual addition
+        """
+        assert (
+            atom_level_embeddings.shape[0] == self.n_conformers
+        ), "Number of conformers must be consistent"
+
+        # Subset to [:atom_level_embedding_dim]
+        if atom_level_embeddings.shape[-1] > self.atom_level_embedding_dim:
+            atom_level_embeddings = atom_level_embeddings[
+                ..., : self.atom_level_embedding_dim
+            ]
+        elif atom_level_embeddings.shape[-1] < self.atom_level_embedding_dim:
+            raise ValueError(
+                f"Atom-level embedding dimension {atom_level_embeddings.shape[-1]} is less than the expected dimension {self.atom_level_embedding_dim}"
+            )
+
+        # Process atom-level embeddings to get shape [n_conformers, n_atom, c_atompair]
+        processed_embeddings: Float[torch.Tensor, "n_conformers n_atom c_atompair"] = (
+            self.process_atom_level_embedding(atom_level_embeddings)
+        )
+
+        # Pad with zeros if we don't have enough conformers
+        current_n_conformers = processed_embeddings.shape[0]
+        if current_n_conformers < self.n_conformers:
+            # Pad with zeros at the beginning
+            padding_size = self.n_conformers - current_n_conformers
+            padding: Float[torch.Tensor, "padding_size n_atom c_atompair"] = (
+                torch.zeros(
+                    padding_size,
+                    processed_embeddings.shape[1],
+                    processed_embeddings.shape[2],
+                    device=processed_embeddings.device,
+                    dtype=processed_embeddings.dtype,
+                )
+            )
+            processed_embeddings = torch.cat([padding, processed_embeddings], dim=0)
+        elif current_n_conformers > self.n_conformers:
+            # Truncate to n_conformers
+            processed_embeddings = processed_embeddings[: self.n_conformers]
+
+        # Reshape from [n_conformers, n_atom, c_atompair] to [n_atom, n_conformers * c_atompair]
+        reshaped_embeddings: Float[torch.Tensor, "n_atom n_conformers*c_atompair"] = (
+            rearrange(processed_embeddings, "c n d -> n (c d)")
+        )
+
+        # Final MLP to get [n_atom, c_atom]
+        result: Float[torch.Tensor, "n_atom c_atom"] = (
+            self.conformers_to_atom_single_embedding(reshaped_embeddings)
+        )
+
+        return result

rf3/model/layers/outer_product.py

@@ -0,0 +1,59 @@
+import torch
+import torch.nn as nn
+from rf3.util_module import init_lecun_normal
+
+from foundry.training.checkpoint import activation_checkpointing
+
+
+class OuterProductMean(nn.Module):
+    def __init__(self, d_msa=256, d_pair=128, d_hidden=16, p_drop=0.15):
+        super(OuterProductMean, self).__init__()
+        self.norm = nn.LayerNorm(d_msa)
+        self.proj_left = nn.Linear(d_msa, d_hidden)
+        self.proj_right = nn.Linear(d_msa, d_hidden)
+        self.proj_out = nn.Linear(d_hidden * d_hidden, d_pair)
+
+        self.reset_parameter()
+
+    def reset_parameter(self):
+        # normal initialization
+        self.proj_left = init_lecun_normal(self.proj_left)
+        self.proj_right = init_lecun_normal(self.proj_right)
+        nn.init.zeros_(self.proj_left.bias)
+        nn.init.zeros_(self.proj_right.bias)
+
+        # zero initialize output
+        nn.init.zeros_(self.proj_out.weight)
+        nn.init.zeros_(self.proj_out.bias)
+
+    def forward(self, msa):
+        B, N, L = msa.shape[:3]
+        msa = self.norm(msa)
+        left = self.proj_left(msa)
+        right = self.proj_right(msa)
+        right = right / float(N)
+        out = torch.einsum("bsli,bsmj->blmij", left, right).reshape(B, L, L, -1)
+        out = self.proj_out(out)
+
+        return out
+
+
+class OuterProductMean_AF3(nn.Module):
+    def __init__(self, c_msa_embed, c_outer_product, c_out):
+        super(OuterProductMean_AF3, self).__init__()
+        self.norm = nn.LayerNorm(c_msa_embed)
+        self.proj_left = nn.Linear(c_msa_embed, c_outer_product)
+        self.proj_right = nn.Linear(c_msa_embed, c_outer_product)
+        self.proj_out = nn.Linear(c_outer_product * c_outer_product, c_out)
+
+    @activation_checkpointing
+    def forward(self, msa):
+        B, N, L = msa.shape[:3]
+        msa = self.norm(msa)
+        left = self.proj_left(msa)
+        right = self.proj_right(msa)
+        right = right / float(N)
+        out = torch.einsum("bsli,bsmj->blmij", left, right).reshape(B, L, L, -1)
+        out = self.proj_out(out)
+
+        return out