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

hippoformer/hippoformer.py

@@ -285,7 +285,7 @@ class Attention(Module):
         q, k1, v1, k2, v2 = tuple(self.split_heads(t) for t in (q, k1, v1, k2, v2))
 
         if exists(kv_cache):
-            ck1, cv1, vk2, cv2 = kv_cache
+            ck1, cv1, ck2, cv2 = kv_cache
             k1 = cat((ck1, k1), dim = -2)
             v1 = cat((cv1, v1), dim = -2)
             k2 = cat((ck2, k2), dim = -2)
@@ -300,12 +300,12 @@ class Attention(Module):
 
             i, j = sim.shape[-2:]
 
-            j_seq = arange(j, device = device)[:, None]
-            i_seq = arange(i, device = device)[None, :] + (j - i)
+            j_seq = arange(i, device = device)[:, None]
+            i_seq = arange(j, device = device)[None, :] + (j - i)
 
             windowed_causal_mask_without_diagonal = (i_seq > j_seq) & ((i_seq - j_seq) <= self.window_size)
 
-            sim = sim.masked_fill(windowed_causal_mask_without_diagonal, -torch.finfo(sim.dtype).max)
+            sim = sim.masked_fill(~windowed_causal_mask_without_diagonal, -torch.finfo(sim.dtype).max)
 
             # attention sink, for token as well as for attention sinking - from gpt-oss
 
@@ -365,7 +365,7 @@ class TEMTransformerBlock(Module):
 
         x = self.ff(x) + x
 
-        next_kv_cache = next_kv_cache[:, -self.window_size:]
+        next_kv_cache = tuple(t[:, -self.window_size:] for t in next_kv_cache)
 
         return x, next_kv_cache
 
@@ -401,7 +401,7 @@ class TEMTransformer(Module):
                 **transformer_kwargs
             )
 
-            layers.append(block)
+            self.layers.append(block)
 
     def forward(
         self,
@@ -411,7 +411,7 @@ class TEMTransformer(Module):
         prev_kv_cache = None  # for the specialized transformer blocks for inducing the grid-cells
     ):
         
-        structure, next_hiddens = self.gru_path_integrator(actions, prev_hiddens)
+        structure, next_hiddens = self.path_integrator(actions, prev_hiddens)
 
         encoded_sensory = self.sensory_encoder(sensory)
 
@@ -484,7 +484,7 @@ class mmTEM(Module):
         self.to_values = nn.Linear(dim_joint_rep, dim, bias = False)
 
         self.to_learned_optim_hparams = nn.Linear(dim_joint_rep, 3, bias = False) # for learning rate, forget gate, and momentum
-        self.assoc_scan = AssocScan(*assoc_scan_kwargs)
+        self.assoc_scan = AssocScan(**assoc_scan_kwargs)
 
         self.meta_memory_mlp = create_mlp(
             dim = dim * 2,
@@ -495,7 +495,7 @@ class mmTEM(Module):
         )
 
         def forward_with_mse_loss(params, keys, values):
-            pred = functional_call(self.meta_memory_mlp, params, keys)
+            pred = functional_call(self.meta_memory_mlp, params, (keys,))
             return F.mse_loss(pred, values)
 
         grad_fn = grad(forward_with_mse_loss)
@@ -596,7 +596,7 @@ class mmTEM(Module):
 
         # 2b. structure from structure
 
-        decoded_structure, decoded_encoded_sensory = self.retrieve(zeros_like(structural_codes), encoded_sensory)
+        decoded_structure, decoded_encoded_sensory = self.retrieve(structural_codes, zeros_like(encoded_sensory))
 
         structure_from_structure_loss = F.mse_loss(decoded_structure, structural_codes)
 

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

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

hippoformer-0.0.16.dist-info/RECORD

@@ -0,0 +1,6 @@
+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,,

{hippoformer-0.0.14.dist-info → hippoformer-0.0.16.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: hatchling 1.27.0
+Generator: hatchling 1.28.0
 Root-Is-Purelib: true
 Tag: py3-none-any

hippoformer-0.0.14.dist-info/RECORD

@@ -1,6 +0,0 @@
-hippoformer/__init__.py,sha256=A7N8GsRAZH4yP-L5hb7IVDnNjnhfjNyolg5MZ6vnGyE,71
-hippoformer/hippoformer.py,sha256=RqHobJTA36mJ-YFutteMx5_6QrrSqI370C1sw5AYrUE,20818
-hippoformer-0.0.14.dist-info/METADATA,sha256=W8k541yoURENT9RFpxZtu7O08r_R8joyzjj3ikwFbH0,3093
-hippoformer-0.0.14.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-hippoformer-0.0.14.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-hippoformer-0.0.14.dist-info/RECORD,,