hippoformer 0.0.16__py3-none-any.whl → 0.0.17__py3-none-any.whl

hippoformer/hippoformer.py

@@ -300,8 +300,8 @@ class Attention(Module):
 
             i, j = sim.shape[-2:]
 
-            j_seq = arange(i, device = device)[:, None]
-            i_seq = arange(j, device = device)[None, :] + (j - i)
+            i_seq = arange(i, device = device)[:, None] + (j - i)
+            j_seq = arange(j, device = device)[None, :]
 
             windowed_causal_mask_without_diagonal = (i_seq > j_seq) & ((i_seq - j_seq) <= self.window_size)
 
@@ -407,18 +407,22 @@ class TEMTransformer(Module):
         self,
         sensory,
         actions,
-        prev_hiddens = None,  # for the GRU based path integrator
-        prev_kv_cache = None  # for the specialized transformer blocks for inducing the grid-cells
+        prev_hiddens = None,   # for the GRU based path integrator
+        prev_kv_cache = None,  # for the specialized transformer blocks for inducing the grid-cells
+        return_memories = False
     ):
         
         structure, next_hiddens = self.path_integrator(actions, prev_hiddens)
 
         encoded_sensory = self.sensory_encoder(sensory)
 
+        prev_kv_cache = default(prev_kv_cache, (None,) * len(self.layers))
+        iter_prev_kv_cache = iter(prev_kv_cache)
+
         next_kv_cache = []
 
         for layer in self.layers:
-            structure, layer_next_cache = layer(structure, encoded_sensory)
+            structure, layer_next_cache = layer(structure, encoded_sensory, kv_cache = next(iter_prev_kv_cache, None))
             next_kv_cache.append(layer_next_cache)
 
         decoded_sensory = self.sensory_decoder(structure)
@@ -427,7 +431,10 @@ class TEMTransformer(Module):
 
         pred_loss = F.mse_loss(encoded_sensory, decoded_sensory)
 
-        return pred_loss
+        if not return_memories:
+            return pred_loss
+
+        return pred_loss, next_memories
 
 # proposed mmTEM
 

{hippoformer-0.0.16.dist-info → hippoformer-0.0.17.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hippoformer
-Version: 0.0.16
+Version: 0.0.17
 Summary: hippoformer
 Project-URL: Homepage, https://pypi.org/project/hippoformer/
 Project-URL: Repository, https://github.com/lucidrains/hippoformer

hippoformer-0.0.17.dist-info/RECORD

@@ -0,0 +1,6 @@
+hippoformer/__init__.py,sha256=A7N8GsRAZH4yP-L5hb7IVDnNjnhfjNyolg5MZ6vnGyE,71
+hippoformer/hippoformer.py,sha256=fWgCak1szOJpoL-6bFuoJk8Z7RGxjlcL7-8kmGYhAfU,21121
+hippoformer-0.0.17.dist-info/METADATA,sha256=XfCPGG4gyxvTSrZzV-H-L5I6WzN7MKhYXG0r7puPp1w,3093
+hippoformer-0.0.17.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+hippoformer-0.0.17.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+hippoformer-0.0.17.dist-info/RECORD,,

hippoformer-0.0.16.dist-info/RECORD

@@ -1,6 +0,0 @@
-hippoformer/__init__.py,sha256=A7N8GsRAZH4yP-L5hb7IVDnNjnhfjNyolg5MZ6vnGyE,71
-hippoformer/hippoformer.py,sha256=CchBXGVmVioT2eBLGp44wvC8flv9e-TQLLovcz9G0ts,20842
-hippoformer-0.0.16.dist-info/METADATA,sha256=L77ovZLV2uxu_AIRXNTEJf6h7K5xy21WGHDanJbVC6M,3093
-hippoformer-0.0.16.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-hippoformer-0.0.16.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-hippoformer-0.0.16.dist-info/RECORD,,