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

titans_pytorch-0.0.15/.github/workflows/test.yaml

@@ -0,0 +1,24 @@
+name: Tests the examples in README
+on: [push, pull_request]
+
+env:
+  TYPECHECK: True
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Install Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+      - name: Install dependencies
+        run: |
+          python -m pip install uv
+          python -m uv pip install --upgrade pip
+          python -m uv pip install torch --index-url https://download.pytorch.org/whl/nightly/cpu
+          python -m uv pip install -e .[test]
+      - name: Test with pytest
+        run: |
+          python -m pytest tests/

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.0.14
+Version: 0.0.15
 Summary: Titans
 Project-URL: Homepage, https://pypi.org/project/titans-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/titans-pytorch
@@ -109,3 +109,13 @@ $ python train.py
     url     = {https://api.semanticscholar.org/CorpusID:275212078}
 }
 ```
+
+```bibtex
+@software{Kyrylov_Accelerated_Scan_2024,
+    author  = {Kyrylov, Volodymyr},
+    doi     = {10.5281/zenodo.10600962},
+    title   = {Accelerated Scan},
+    version = {0.1.2},
+    year    = {2024}
+}
+```

{titans_pytorch-0.0.14 → titans_pytorch-0.0.15}/README.md

@@ -58,3 +58,13 @@ $ python train.py
     url     = {https://api.semanticscholar.org/CorpusID:275212078}
 }
 ```
+
+```bibtex
+@software{Kyrylov_Accelerated_Scan_2024,
+    author  = {Kyrylov, Volodymyr},
+    doi     = {10.5281/zenodo.10600962},
+    title   = {Accelerated Scan},
+    version = {0.1.2},
+    year    = {2024}
+}
+```

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

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

titans_pytorch-0.0.15/tests/test_titans.py

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

titans_pytorch-0.0.15/titans_pytorch/__init__.py

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

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

@@ -6,7 +6,7 @@ import torch
 from torch import nn, Tensor
 import torch.nn.functional as F
 from torch.nn import Linear, Module
-from torch.func import functional_call, vmap, grad_and_value
+from torch.func import functional_call, vmap, grad
 
 from tensordict import TensorDict
 
@@ -57,7 +57,7 @@ def pack_one_with_inverse(t, pattern):
 
 # classes
 
-class MLP(Module):
+class MemoryMLP(Module):
     def __init__(
         self,
         dim,
@@ -80,6 +80,35 @@ class MLP(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):
@@ -122,7 +151,7 @@ class NeuralMemory(Module):
         # memory mlp
 
         if not exists(model):
-            model = MLP(dim_head, **default_mlp_kwargs)
+            model = MemoryMLP(dim_head, **default_mlp_kwargs)
 
         assert not exists(next(model.buffers(), None)), 'model cannot have buffers for now'
 
@@ -141,7 +170,7 @@ class NeuralMemory(Module):
             loss = self.store_memory_loss_fn(pred, target) # simple mse loss in paper - eq (12) - |M(k) - v|²
             return loss
 
-        self.per_sample_grad_and_value_fn = vmap(grad_and_value(forward_and_loss), in_dims = (None, 0, 0))
+        self.per_sample_grad_fn = vmap(grad(forward_and_loss), in_dims = (None, 0, 0))
 
         # queries for retrieving from the model
 
@@ -235,7 +264,7 @@ class NeuralMemory(Module):
 
         # get grads and extra auxiliary loss (for backwarding through qkv projection in base neural memory module)
 
-        grads, aux_store_loss = self.per_sample_grad_and_value_fn(dict(curr_weights), keys, values)
+        grads = self.per_sample_grad_fn(dict(curr_weights), keys, values)
 
         grads = TensorDict(grads)
 
@@ -305,7 +334,7 @@ class NeuralMemory(Module):
 
         next_state = (curr_weights + last_update, next_momentum)
 
-        return updates, next_state, aux_store_loss.mean() / chunk_size
+        return updates, next_state
 
     def retrieve_memories(
         self,
@@ -382,6 +411,7 @@ class NeuralMemory(Module):
     def forward(
         self,
         seq,
+        store_seq = None,
         past_state: tuple[dict[str, Tensor], dict[str, Tensor]] | None = None,
         return_next_memories = False
     ):
@@ -396,7 +426,9 @@ class NeuralMemory(Module):
         if not exists(past_state):
             past_state = self.init_weights_and_momentum()
 
-        updates, next_memories, aux_kv_mse_loss = self.store_memories(seq, past_state)
+        store_seq = default(store_seq, seq)
+
+        updates, next_memories = self.store_memories(store_seq, past_state)
 
         past_weights, _ = past_state
 
@@ -405,4 +437,4 @@ class NeuralMemory(Module):
         if not return_next_memories:
             return retrieved
 
-        return retrieved, next_memories, aux_kv_mse_loss
+        return retrieved, next_memories

titans_pytorch-0.0.14/titans_pytorch/__init__.py

@@ -1,3 +0,0 @@
-from titans_pytorch.titans import (
-    NeuralMemory
-)