complete basic prelayernorm + adaptive layernorm as well as condition…

… wrapper for both attention and feedforward with the adaln-zero

alphafold3_pytorch/__init__.py

@@ -4,13 +4,19 @@
 )
 
 from alphafold3_pytorch.alphafold3 import (
+    PreLayerNorm,
+    AdaptiveLayerNorm,
+    ConditionWrapper,
     Transition,
     Alphafold3
 )
 
 __all__ = [
     Attention,
     Attend,
+    PreLayerNorm,
+    AdaptiveLayerNorm,
+    ConditionWrapper,
     Transition,
     Alphafold3
 ]

alphafold3_pytorch/alphafold3.py

@@ -1,6 +1,9 @@
+from __future__ import annotations
 from functools import partial
 
 import torch
+from torch import nn
+from torch import Tensor
 import torch.nn.functional as F
 from torch.nn import Module, Linear, Sequential
 
@@ -11,6 +14,8 @@
     typecheck
 )
 
+from alphafold3_pytorch.attention import Attention
+
 # constants
 
 LinearNoBias = partial(Linear, bias = False)
@@ -61,6 +66,106 @@ def forward(
 
         return self.ff(x)
 
+# normalization
+# both pre layernorm as well as adaptive layernorm wrappers
+
+class PreLayerNorm(Module):
+    @typecheck
+    def __init__(
+        self,
+        fn: Attention | Transition,
+        *,
+        dim,
+    ):
+        super().__init__()
+        self.fn = fn
+        self.norm = nn.LayerNorm(dim)
+
+    @typecheck
+    def forward(
+        self,
+        x: Tensor,
+        **kwargs
+    ):
+        x = self.norm(x)
+        return self.fn(x, **kwargs)
+
+class AdaptiveLayerNorm(Module):
+    """ Algorithm 26 """
+
+    def __init__(
+        self,
+        *,
+        dim,
+        dim_cond
+    ):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim, elementwise_affine = False)
+        self.norm_cond = nn.LayerNorm(dim_cond, bias = False)
+
+        self.to_gamma = nn.Sequential(
+            nn.Linear(dim_cond, dim),
+            nn.Sigmoid()
+        )
+
+        self.to_beta = nn.Linear(dim_cond, dim, bias = False)
+
+    @typecheck
+    def forward(
+        self,
+        x: Tensor,
+        cond: Tensor
+    ):
+        normed = self.norm(x)
+        normed_cond = self.norm_cond(cond)
+
+        gamma = self.to_gamma(normed_cond)
+        beta = self.to_beta(normed_cond)
+        return normed * gamma + beta
+
+class ConditionWrapper(Module):
+    """ Algorithm 25 """
+
+    @typecheck
+    def __init__(
+        self,
+        fn: Attention | Transition,
+        *,
+        dim,
+        dim_cond,
+        adaln_zero_bias_init_value = -2.
+    ):
+        super().__init__()
+        self.fn = fn
+        self.adaptive_norm = AdaptiveLayerNorm(dim = dim, dim_cond = dim_cond)
+
+        adaln_zero_gamma_linear = nn.Linear(dim_cond, dim)
+        nn.init.zeros_(adaln_zero_gamma_linear.weight)
+        nn.init.constant_(adaln_zero_gamma_linear.bias, adaln_zero_bias_init_value)
+
+        self.to_adaln_zero_gamma = nn.Sequential(
+            adaln_zero_gamma_linear,
+            nn.Sigmoid()
+        )
+
+        self.to_adaln_zero_beta = nn.Linear(dim_cond, dim, bias = False)
+
+    @typecheck
+    def forward(
+        self,
+        x: Tensor,
+        *,
+        cond: Tensor,
+        **kwargs
+    ):
+        x = self.adaptive_norm(x, cond = cond)
+
+        out = self.fn(x, **kwargs)
+
+        gamma = self.to_adaln_zero_gamma(cond)
+        beta = self.to_adaln_zero_beta(cond)
+        return out * gamma + beta
+
 # main class
 
 class Alphafold3(Module):