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

titans_pytorch-0.0.16/.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.16}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.0.14
+Version: 0.0.16
 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.16}/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.16}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "titans-pytorch"
-version = "0.0.14"
+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.14 → titans_pytorch-0.0.16}/titans_pytorch/__init__.py

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

{titans_pytorch-0.0.14 → titans_pytorch-0.0.16}/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,
@@ -122,7 +122,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 +141,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 +235,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 +305,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 +382,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 +397,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 +408,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-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