titans-pytorch 0.0.15__tar.gz → 0.0.16__tar.gz

{titans_pytorch-0.0.15 → titans_pytorch-0.0.16}/PKG-INFO

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

{titans_pytorch-0.0.15 → titans_pytorch-0.0.16}/pyproject.toml

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

titans_pytorch-0.0.16/tests/test_titans.py

@@ -0,0 +1,28 @@
+import torch
+import pytest
+
+def test_titans():
+    from titans_pytorch import NeuralMemory
+
+    mem = NeuralMemory(
+        dim = 384,
+        chunk_size = 64,
+    )
+
+    seq = torch.randn(2, 1024, 384)
+    retrieved = mem(seq)
+
+    assert seq.shape == retrieved.shape
+
+def test_titans_attn_memory():
+    from titans_pytorch.titans_attn_memory import NeuralMemory
+
+    mem = NeuralMemory(
+        dim = 384,
+        chunk_size = 64,
+    )
+
+    seq = torch.randn(2, 1024, 384)
+    retrieved = mem(seq)
+
+    assert seq.shape == retrieved.shape

{titans_pytorch-0.0.15 → titans_pytorch-0.0.16}/titans_pytorch/__init__.py

@@ -1,5 +1,4 @@
 from titans_pytorch.titans import (
     NeuralMemory,
     MemoryMLP,
-    MemoryAttention
 )

{titans_pytorch-0.0.15 → titans_pytorch-0.0.16}/titans_pytorch/titans.py

@@ -80,35 +80,6 @@ class MemoryMLP(Module):
 
         return x
 
-# improvised attention as memory module
-# todo - expand if see signal in experiments
-
-class MemoryAttention(Module):
-    def __init__(
-        self,
-        dim
-    ):
-        super().__init__()
-        self.weights = nn.ParameterList([
-            nn.Parameter(torch.randn(dim, dim)), # queries
-            nn.Parameter(torch.randn(dim, dim)), # keys
-            nn.Parameter(torch.randn(dim, dim)), # values
-        ])
-
-    def forward(self, x):
-        wq, wk, wv = self.weights
-
-        q = x @ wq
-        k = x @ wk
-        v = x @ wv
-
-        sim = q @ k.transpose(-1, -2)
-
-        attn = sim.softmax(dim = -1)
-
-        out = attn @ v
-        return out
-
 # main neural memory
 
 def default_loss_fn(pred, target):

{titans_pytorch-0.0.15 → titans_pytorch-0.0.16}/train.py

@@ -13,7 +13,11 @@ from local_attention import LocalTransformer
 
 from taylor_series_linear_attention import TaylorSeriesLinearAttn
 
-from titans_pytorch.titans import NeuralMemory
+from titans_pytorch.titans import (
+    NeuralMemory,
+    MemoryAttention,
+    MemoryMLP
+)
 
 # constants
 

titans_pytorch-0.0.15/tests/test_titans.py

@@ -1,15 +0,0 @@
-import pytest
-
-import torch
-from titans_pytorch import NeuralMemory
-
-def test_titans():
-    mem = NeuralMemory(
-        dim = 384,
-        chunk_size = 64,
-    )
-
-    seq = torch.randn(2, 1024, 384)
-    retrieved = mem(seq)
-
-    assert seq.shape == retrieved.shape