alphagenome-pytorch 0.0.4__tar.gz → 0.0.5__tar.gz

{alphagenome_pytorch-0.0.4 → alphagenome_pytorch-0.0.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: alphagenome-pytorch
-Version: 0.0.4
+Version: 0.0.5
 Summary: AlphaGenome
 Project-URL: Homepage, https://pypi.org/project/alphagenome-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/alphagenome

{alphagenome_pytorch-0.0.4 → alphagenome_pytorch-0.0.5}/alphagenome_pytorch/alphagenome.py

@@ -15,6 +15,7 @@ from einops import rearrange, repeat, einsum
 # b - batch
 # h - heads
 # n - sequence
+# p - relative positions
 # d - feature dimension
 
 # constants
@@ -79,9 +80,15 @@ def apply_rotary_pos_emb(pos, t):
 # 'central mask features' - relative positions for constituting pairwise rep
 
 class RelativePosFeatures(Module):
+    def __init__(self, pool_size = 16):
+        super().__init__()
+        self.pool_size = pool_size
+
     def forward(self, single):
 
         _, seq_len, dim = single.shape
+
+        seq_len //= self.pool_size
         half_dim = dim // 2
 
         rel_pos = arange(2 * seq_len - 1) - (seq_len - 1)
@@ -271,12 +278,17 @@ class SingleToPairwise(Module):
 
             q_rel_bias, k_rel_bias = self.qk_rel_pos_bias
 
-            rel_q = relative_shift(einsum(q + q_rel_bias, rel_pos_encoding, 'b i h d, j h d -> b i j h'))
-            rel_k = relative_shift(einsum(k + k_rel_bias, rel_pos_encoding, 'b i h d, j h d -> b i j h'))
+            rel_q = relative_shift(einsum(q + q_rel_bias, rel_pos_encoding, 'b n h d, p h d -> b n p h'))
+            rel_k = relative_shift(einsum(k + k_rel_bias, rel_pos_encoding, 'b n h d, p h d -> b n p h'))
 
-            rel_sim = einsum(rel_q, rel_k, 'b i p h, b j p h -> b i j h')
+            _, seq, rel_pos_seq, _ = rel_q.shape
+            crop_padding = (rel_pos_seq - seq) // 2
 
-            sim = (sim + rel_sim) * 0.5
+            rel_q, rel_k = tuple(t[..., crop_padding:(crop_padding + seq), :] for t in (rel_q, rel_k))
+
+            rel_sim = add('b i j d, b j i d -> b i j d', rel_q, rel_k) * 0.5
+
+            sim = sim + rel_sim
 
         pairwise_from_sim = self.qk_to_pairwise(sim)
 
@@ -353,6 +365,7 @@ class TransformerTower(Module):
         dim_pairwise = None,
         pairwise_every_num_single_blocks = 2,   # how often to do a pairwise block
         single_to_pairwise_heads = 32,          # they did 32
+        pool_size = 16,
         attn_kwargs: dict = dict(),
         ff_kwargs: dict = dict()
     ):
@@ -363,7 +376,7 @@ class TransformerTower(Module):
 
         self.pairwise_every = pairwise_every_num_single_blocks
 
-        self.rel_pos_features = RelativePosFeatures()
+        self.rel_pos_features = RelativePosFeatures(pool_size)
 
         self.rotary_emb = RotaryEmbedding(dim_head_qk, max_positions = max_positions)
 
@@ -381,7 +394,7 @@ class TransformerTower(Module):
             single_to_pairwise, pairwise_attn, pairwise_ff = None, None, None
 
             if divisible_by(layer_index, self.pairwise_every):
-                single_to_pairwise = SingleToPairwise(dim = dim, dim_pairwise = dim_pairwise, heads = single_to_pairwise_heads)
+                single_to_pairwise = SingleToPairwise(dim = dim, dim_pairwise = dim_pairwise, heads = single_to_pairwise_heads, pool_size = pool_size)
                 pairwise_attn = PairwiseRowAttention(dim_pairwise)
                 pairwise_ff = FeedForward(dim = dim_pairwise, expansion_factor = ff_expansion_factor)
 

{alphagenome_pytorch-0.0.4 → alphagenome_pytorch-0.0.5}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "alphagenome-pytorch"
-version = "0.0.4"
+version = "0.0.5"
 description = "AlphaGenome"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }