hippoformer 0.0.7__tar.gz → 0.0.8__tar.gz

{hippoformer-0.0.7 → hippoformer-0.0.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hippoformer
-Version: 0.0.7
+Version: 0.0.8
 Summary: hippoformer
 Project-URL: Homepage, https://pypi.org/project/hippoformer/
 Project-URL: Repository, https://github.com/lucidrains/hippoformer
@@ -35,6 +35,7 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.9
 Requires-Dist: assoc-scan
+Requires-Dist: beartype
 Requires-Dist: einops>=0.8.1
 Requires-Dist: einx>=0.3.0
 Requires-Dist: torch>=2.4
@@ -50,8 +51,6 @@ Description-Content-Type: text/markdown
 
 Implementation of [Hippoformer](https://openreview.net/forum?id=hxwV5EubAw), Integrating Hippocampus-inspired Spatial Memory with Transformers
 
-[Temporary Discord](https://discord.gg/MkACrrkrYR)
-
 ## Citations
 
 ```bibtex

{hippoformer-0.0.7 → hippoformer-0.0.8}/README.md

@@ -4,8 +4,6 @@
 
 Implementation of [Hippoformer](https://openreview.net/forum?id=hxwV5EubAw), Integrating Hippocampus-inspired Spatial Memory with Transformers
 
-[Temporary Discord](https://discord.gg/MkACrrkrYR)
-
 ## Citations
 
 ```bibtex

{hippoformer-0.0.7 → hippoformer-0.0.8}/hippoformer/hippoformer.py

@@ -7,6 +7,8 @@ from torch.nn import Module
 from torch.jit import ScriptModule, script_method
 from torch.func import vmap, grad, functional_call
 
+from beartype import beartype
+
 from einx import multiply
 from einops import repeat, rearrange, pack, unpack
 from einops.layers.torch import Rearrange
@@ -36,6 +38,80 @@ def pack_with_inverse(t, pattern):
 def l2norm(t):
     return F.normalize(t, dim = -1)
 
+# sensory encoder decoder for 2d
+
+grid_sensory_enc_dec = (
+    create_mlp(
+        dim = 32 * 2,
+        dim_in = 9,
+        dim_out = 32,
+        depth = 3,
+    ),
+    create_mlp(
+        dim = 32 * 2,
+        dim_in = 32,
+        dim_out = 9,
+        depth = 3,
+    ),
+)
+
+# sensory encoder decoder for 3d maze
+
+class EncoderPackTime(Module):
+    def __init__(self, fn: Module):
+        super().__init__()
+        self.fn = fn
+
+    def forward(self, x):
+        x = rearrange(x, 'b c t h w -> b t c h w')
+        x, packed_shape = pack([x], '* c h w')
+
+        x = self.fn(x)
+
+        x, = unpack(x, packed_shape, '* d')
+        print(x.shape)
+        return x
+
+class DecoderPackTime(Module):
+    def __init__(self, fn: Module):
+        super().__init__()
+        self.fn = fn
+
+    def forward(self, x):
+        x, packed_shape = pack(x, '* d')
+
+        x = self.fn(x)
+
+        x = unpack(x, packed_shape, '* c h w')
+        x = rearrange(x, 'b t c h w -> b c t h w')
+        return x
+
+maze_sensory_enc_dec = (
+    EncoderPackTime(nn.Sequential(
+        nn.Conv2d(3, 16, 7, 2, padding = 3),
+        nn.ReLU(),
+        nn.Conv2d(16, 32, 3, 2, 1),
+        nn.ReLU(),
+        nn.Conv2d(32, 64, 3, 2, 1),
+        nn.ReLU(),
+        nn.Conv2d(64, 128, 3, 2, 1),
+        nn.ReLU(),
+        Rearrange('b ... -> b (...)'),
+        nn.Linear(2048, 32)
+    )),
+    DecoderPackTime(nn.Sequential(
+        nn.Linear(32, 2048),
+        Rearrange('b (c h w) -> b c h w', c = 128, h = 4),
+        nn.ConvTranspose2d(128, 64, 3, 2, 1, output_padding = (1, 1)),
+        nn.ReLU(),
+        nn.ConvTranspose2d(64, 32, 3, 2, 1, output_padding = (1, 1)),
+        nn.ReLU(),
+        nn.ConvTranspose2d(32, 16, 3, 2, 1, output_padding = (1, 1)),
+        nn.ReLU(),
+        nn.ConvTranspose2d(16, 3, 3, 2, 1, output_padding = (1, 1))
+    ))
+)
+
 # path integration
 
 class RNN(ScriptModule):
@@ -114,12 +190,12 @@ class PathIntegration(Module):
 # proposed mmTEM
 
 class mmTEM(Module):
+    @beartype
     def __init__(
         self,
         dim,
         *,
-        sensory_encoder: Module,
-        sensory_decoder: Module,
+        sensory_encoder_decoder: tuple[Module, Module],
         dim_sensory,
         dim_action,
         dim_encoded_sensory,
@@ -139,6 +215,8 @@ class mmTEM(Module):
 
         # sensory
 
+        sensory_encoder, sensory_decoder = sensory_encoder_decoder
+
         self.sensory_encoder = sensory_encoder
         self.sensory_decoder = sensory_decoder
 

{hippoformer-0.0.7 → hippoformer-0.0.8}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "hippoformer"
-version = "0.0.7"
+version = "0.0.8"
 description = "hippoformer"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -25,6 +25,7 @@ classifiers=[
 
 dependencies = [
     "assoc-scan",
+    "beartype",
     "einx>=0.3.0",
     "einops>=0.8.1",
     "torch>=2.4",

{hippoformer-0.0.7 → hippoformer-0.0.8}/tests/test_hippoformer.py

@@ -1,4 +1,5 @@
 import pytest
+param = pytest.mark.parametrize
 
 import torch
 
@@ -14,16 +15,40 @@ def test_path_integrate():
 
     assert structure_codes.shape == (2, 16, 64)
 
-def test_mm_tem():
+@param('sensory_type', ('naive', '2d', '3d'))
+def test_mm_tem(
+    sensory_type
+):
     import torch
     from hippoformer.hippoformer import mmTEM
 
     from torch.nn import Linear
 
+    if sensory_type == 'naive':
+        enc_dec = (
+            Linear(11, 32),
+            Linear(32, 11)
+        )
+        sensory = torch.randn(2, 16, 11)
+
+    elif sensory_type == '2d':
+
+        from hippoformer.hippoformer import grid_sensory_enc_dec
+
+        enc_dec = grid_sensory_enc_dec
+        sensory = torch.randn(2, 16, 9)
+
+    elif sensory_type == '3d':
+
+        from hippoformer.hippoformer import maze_sensory_enc_dec
+
+        enc_dec = maze_sensory_enc_dec
+
+        sensory = torch.randn(2, 3, 16, 64, 64)
+
     model = mmTEM(
         dim = 32,
-        sensory_encoder = Linear(11, 32),
-        sensory_decoder = Linear(32, 11),
+        sensory_encoder_decoder = enc_dec,
         dim_sensory = 11,
         dim_action = 7,
         dim_structure = 32,
@@ -31,10 +56,6 @@ def test_mm_tem():
     )
 
     actions = torch.randn(2, 16, 7)
-    sensory = torch.randn(2, 16, 11)
-
-    actions = torch.randn(2, 16, 7)
-    sensory = torch.randn(2, 16, 11)
 
     next_params = model(sensory, actions, return_memory_mlp_params = True)
     next_params = model(sensory, actions, memory_mlp_params = next_params, return_memory_mlp_params = True)