alphagenome-pytorch 0.0.5__py3-none-any.whl → 0.0.7__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
 
@@ -267,6 +268,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))
 
@@ -278,15 +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'))
-
-            _, seq, rel_pos_seq, _ = rel_q.shape
-            crop_padding = (rel_pos_seq - seq) // 2
-
-            rel_q, rel_k = tuple(t[..., crop_padding:(crop_padding + seq), :] for t in (rel_q, rel_k))
+            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_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
 

{alphagenome_pytorch-0.0.5.dist-info → alphagenome_pytorch-0.0.7.dist-info}/METADATA

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

alphagenome_pytorch-0.0.7.dist-info/RECORD

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

alphagenome_pytorch-0.0.5.dist-info/RECORD

@@ -1,6 +0,0 @@
-alphagenome_pytorch/__init__.py,sha256=7plC_YRm0UapNCl9hJEhFxXE-ELGKVy-DtuO5GUQxGI,101
-alphagenome_pytorch/alphagenome.py,sha256=5uxmx-x-6KSu-2aUF4AxM-rYSxXuwnr-0meTNhu2vUM,14086
-alphagenome_pytorch-0.0.5.dist-info/METADATA,sha256=0W6pkI9-RvNCOixhC1ius8a4-cwpTx7gsA370FLmnsY,3386
-alphagenome_pytorch-0.0.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-alphagenome_pytorch-0.0.5.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-alphagenome_pytorch-0.0.5.dist-info/RECORD,,