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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: alphagenome-pytorch
-Version: 0.0.4
+Version: 0.0.6
 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.6}/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)
@@ -260,6 +267,8 @@ class SingleToPairwise(Module):
 
         single = self.avg_pool(single)
 
+        pool_seq_len = single.shape[1]
+
         q, k = self.to_qk(single).chunk(2, dim = -1)
         q, k = tuple(self.split_heads(t) for t in (q, k))
 
@@ -271,12 +280,14 @@ 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_q, rel_k = tuple(t[..., :pool_seq_len, :] for t in (rel_q, rel_k))
 
-            rel_sim = einsum(rel_q, rel_k, 'b i p h, b j p h -> b i j h')
+            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) * 0.5
+            sim = sim + rel_sim
 
         pairwise_from_sim = self.qk_to_pairwise(sim)
 
@@ -353,6 +364,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 +375,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 +393,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.6}/pyproject.toml

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