alphagenome-pytorch 0.0.6__py3-none-any.whl → 0.0.8__py3-none-any.whl

alphagenome_pytorch/alphagenome.py

@@ -46,7 +46,8 @@ def relative_shift(t):
     *leading_dims, seq_len, dim = t.shape
     t = F.pad(t, (1, 0), value = 0.)
     t = t.reshape(*leading_dims, dim + 1, seq_len)
-    return t[..., 1:, :].reshape(*leading_dims, seq_len, dim)
+    t = t[..., 1:, :].reshape(*leading_dims, seq_len, dim)
+    return t[..., :, :seq_len]
 
 # rotary, but with attenuation of short relative distance frequencies
 
@@ -280,12 +281,10 @@ 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 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 = relative_shift(einsum(q + q_rel_bias, rel_pos_encoding, 'b n h d, p h d -> b h n p'))
+            rel_k = relative_shift(einsum(k + k_rel_bias, rel_pos_encoding, 'b n h d, p h d -> b h n p'))
 
-            rel_q, rel_k = tuple(t[..., :pool_seq_len, :] 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
+            rel_sim = add('b h i j, b h j i -> b h i j', rel_q, rel_k) * 0.5
 
             sim = sim + rel_sim
 
@@ -435,14 +434,14 @@ class TransformerTower(Module):
             maybe_pairwise_ff
         ) in self.layers:
 
-            single = attn(single, rotary_emb = rotary_emb, pairwise = None) + single
-            single = ff(single) + single
-
             if exists(maybe_single_to_pair):
                 pairwise = maybe_single_to_pair(single, rel_pos_feats) + default(pairwise, 0.)
                 pairwise = maybe_pairwise_attn(pairwise) + pairwise
                 pairwise = maybe_pairwise_ff(pairwise) + pairwise
 
+            single = attn(single, rotary_emb = rotary_emb, pairwise = None) + single
+            single = ff(single) + single
+
         return single, pairwise
 
 # embedding

{alphagenome_pytorch-0.0.6.dist-info → alphagenome_pytorch-0.0.8.dist-info}/METADATA

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

alphagenome_pytorch-0.0.8.dist-info/RECORD

@@ -0,0 +1,6 @@
+alphagenome_pytorch/__init__.py,sha256=7plC_YRm0UapNCl9hJEhFxXE-ELGKVy-DtuO5GUQxGI,101
+alphagenome_pytorch/alphagenome.py,sha256=UjlOJGTMshcaNqmY0r6IRCgRunu7BTZyhbk2vNi5Mis,13948
+alphagenome_pytorch-0.0.8.dist-info/METADATA,sha256=mhlnCRy7Ovq_Gt9ul_iIqws4zdoOcLjfQTBrlbcNic8,3386
+alphagenome_pytorch-0.0.8.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+alphagenome_pytorch-0.0.8.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+alphagenome_pytorch-0.0.8.dist-info/RECORD,,

alphagenome_pytorch-0.0.6.dist-info/RECORD

@@ -1,6 +0,0 @@
-alphagenome_pytorch/__init__.py,sha256=7plC_YRm0UapNCl9hJEhFxXE-ELGKVy-DtuO5GUQxGI,101
-alphagenome_pytorch/alphagenome.py,sha256=f0XX7VnbOABwT4WOadQQ64VgYS0-jsqBconX8qaa1eM,14004
-alphagenome_pytorch-0.0.6.dist-info/METADATA,sha256=kI5ZBXu_SYAQmvXrA6L-QBGS7Bd3Xv2SdPoJ9rnCyjw,3386
-alphagenome_pytorch-0.0.6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-alphagenome_pytorch-0.0.6.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-alphagenome_pytorch-0.0.6.dist-info/RECORD,,