hippoformer 0.0.16__tar.gz → 0.0.17__tar.gz

{hippoformer-0.0.16 → hippoformer-0.0.17}/PKG-INFO

@@ -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.16 → hippoformer-0.0.17}/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 → hippoformer-0.0.17}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "hippoformer"
-version = "0.0.16"
+version = "0.0.17"
 description = "hippoformer"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }