File size: 2,216 Bytes
8e04e6f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 | import torch
from .adaptive_ln_scale import (
AdaptiveLayerNorm,
AdaptiveOutputScale,
)
from .swiglu import SwiGLU
# Code from Lucidrain's implementation of AF3
# https://github.com/lucidrains/alphafold3-pytorch
class Transition(torch.nn.Module):
"""Transition layer."""
def __init__(self, dim, expansion_factor=4, layer_norm=False):
super().__init__()
dim_inner = int(dim * expansion_factor)
self.use_layer_norm = layer_norm
if self.use_layer_norm:
self.ln = torch.nn.LayerNorm(dim)
self.swish_linear = torch.nn.Sequential(
torch.nn.Linear(dim, dim_inner * 2, bias=False),
SwiGLU(),
)
self.linear_out = torch.nn.Linear(dim_inner, dim, bias=False)
def forward(self, x, mask):
"""
Args:
x: Input sequence representation, shape [b, n, dim]
mask: binary, shape [b, n]
Returns:
Updated sequence representation, shape [b, n, dim]
"""
if self.use_layer_norm:
x = self.ln(x)
x = self.linear_out(self.swish_linear(x))
return x * mask[..., None]
class TransitionADALN(torch.nn.Module):
"""Transition layer with adaptive layer norm applied to input and adaptive
scaling aplied to output."""
def __init__(self, *, dim, dim_cond, expansion_factor=4):
super().__init__()
self.adaln = AdaptiveLayerNorm(dim=dim, dim_cond=dim_cond)
self.transition = Transition(
dim=dim, expansion_factor=expansion_factor, layer_norm=False
)
self.scale_output = AdaptiveOutputScale(dim=dim, dim_cond=dim_cond)
def forward(self, x, cond, mask):
"""
Args:
x: Input sequence representation, shape [b, n, dim]
cond: conditioning variables, shape [b, n, dim_cond]
mask: binary mask, shape [b, n]
Returns:
Updated sequence representation, shape [b, n, dim]
"""
x = self.adaln(x, cond, mask) # [b, n, dim]
x = self.transition(x, mask) # [b, n, dim]
x = self.scale_output(x, cond, mask) # [b, n, dim]
return x * mask[..., None] # [b, n, dim]
|