optimize for clarity

flamingo_pytorch/flamingo_pytorch.py

@@ -53,13 +53,10 @@ def forward(self, x, latents):
         q = self.to_q(latents)
 
         # the paper differs from Perceiver in which they also concat the key / values derived from the latents to be attended to
-        lk, lv = self.to_kv(latents).chunk(2, dim = -1)
-        k, v = self.to_kv(x).chunk(2, dim = -1)
+        kv_input = torch.cat((x, latents), dim = -2)
+        k, v = self.to_kv(kv_input).chunk(2, dim = -1)
 
-        k, v, lk, lv, q = rearrange_many((k, v, lk, lv, q), 'b t n (h d) -> b h t n d', h = h)
-
-        k = torch.cat((k, lk), dim = -2)
-        v = torch.cat((v, lv), dim = -2)
+        q, k, v = rearrange_many((q, k, v), 'b t n (h d) -> b h t n d', h = h)
 
         q = q * self.scale
 

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'flamingo-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.0.11',
+  version = '0.0.12',
   license='MIT',
   description = 'Flamingo - Pytorch',
   author = 'Phil Wang',