titans-pytorch 0.1.35__py3-none-any.whl → 0.1.37__py3-none-any.whl

titans_pytorch/__init__.py

@@ -1,5 +1,8 @@
 from titans_pytorch.neural_memory import (
     NeuralMemory,
+)
+
+from titans_pytorch.memory_models import (
     MemoryMLP,
     MemoryAttention,
     FactorizedMemoryMLP,

titans_pytorch/memory_models.py

@@ -0,0 +1,153 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch.nn import Module, ModuleList, Parameter, ParameterList
+
+class MemoryMLP(Module):
+    def __init__(
+        self,
+        dim,
+        depth
+    ):
+        super().__init__()
+        self.weights = ParameterList([Parameter(torch.randn(dim, dim)) for _ in range(depth)])
+
+        for weight in self.weights:
+            nn.init.xavier_uniform_(weight)
+
+    def forward(
+        self,
+        x
+    ):
+        for ind, weight in enumerate(self.weights):
+            is_first = ind == 0
+
+            if not is_first:
+                x = F.silu(x)
+
+            x = x @ weight
+
+        return x
+
+# memory mlp, but with gated residual + final projection
+
+class GatedResidualMemoryMLP(Module):
+    def __init__(
+        self,
+        dim,
+        depth,
+        expansion_factor = 2.
+    ):
+        super().__init__()
+        dim_hidden = int(dim * expansion_factor)
+
+        self.weights = ParameterList([
+            ParameterList([
+                Parameter(torch.randn(dim, dim_hidden)),
+                Parameter(torch.randn(dim_hidden, dim)),
+                Parameter(torch.randn(dim * 2, dim)),
+            ]) for _ in range(depth)
+        ])
+
+        self.final_proj = Parameter(torch.randn(dim, dim))
+
+        for param in self.parameters():
+            nn.init.xavier_uniform_(param)
+
+    def forward(
+        self,
+        x
+    ):
+        for weight1, weight2, to_gates in self.weights:
+            res = x
+
+            hidden = x @ weight1
+            hidden = F.silu(hidden)
+            branch_out = hidden @ weight2
+
+            # gated residual
+
+            gates = cat((branch_out, res), dim = -1) @ to_gates
+            x = res.lerp(branch_out, gates.sigmoid())
+
+        return x @ self.final_proj
+
+# memory mlp with factorized weights
+# so can tradeoff capacity for smaller chunk sizes
+
+class FactorizedMemoryMLP(Module):
+    def __init__(
+        self,
+        dim,
+        depth,
+        k = 32
+    ):
+        super().__init__()
+        self.weights = ParameterList([
+            ParameterList([
+                Parameter(torch.randn(dim, k)),
+                Parameter(torch.randn(k, dim)),
+            ]) for _ in range(depth)
+        ])
+
+        for weight1, weight2 in self.weights:
+            nn.init.xavier_uniform_(weight1)
+            nn.init.xavier_uniform_(weight2)
+
+    def forward(
+        self,
+        x
+    ):
+        for ind, (weight1, weight2) in enumerate(self.weights):
+            is_first = ind == 0
+
+            if not is_first:
+                x = F.silu(x)
+
+            x = x @ weight1 @ weight2
+
+        return x
+
+# improvised attention as memory module
+
+class MemoryAttention(Module):
+    def __init__(
+        self,
+        dim,
+        scale = 8.,
+        expansion_factor = 2.
+    ):
+        super().__init__()
+        self.scale = scale
+        dim_ff_hidden = int(dim * expansion_factor)
+
+        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
+            nn.Parameter(torch.randn(dim, dim_ff_hidden)), # ff w1
+            nn.Parameter(torch.randn(dim_ff_hidden, dim)), # ff w2
+        ])
+
+        for weight in self.weights:
+            nn.init.xavier_uniform_(weight)
+
+    def forward(self, x):
+        wq, wk, wv, ffw1, ffw2 = self.weights
+
+        q = F.normalize(x @ wq, dim = -1)
+        k = F.normalize(x @ wk, dim = -1)
+        v = x @ wv
+
+        attn_out = F.scaled_dot_product_attention(
+            q, k, v,
+            scale = self.scale,
+            is_causal = True
+        )
+
+        x = x + attn_out
+
+        h = F.silu(x @ ffw1)
+        out = h @ ffw2
+
+        return out

titans_pytorch/neural_memory.py

@@ -19,6 +19,10 @@ from titans_pytorch.associative_scan import (
     pad_at_dim
 )
 
+from titans_pytorch.memory_models import(
+    MemoryMLP
+)
+
 import einx
 from einops import rearrange, repeat, reduce, pack, unpack
 from einops.layers.torch import Rearrange, Reduce
@@ -169,157 +173,6 @@ class AttentionPool(Module):
 
         return reduce(x * attn, 'b n c d -> b n d', 'sum')
 
-# classes
-
-class MemoryMLP(Module):
-    def __init__(
-        self,
-        dim,
-        depth
-    ):
-        super().__init__()
-        self.weights = ParameterList([Parameter(torch.randn(dim, dim)) for _ in range(depth)])
-
-        for weight in self.weights:
-            nn.init.xavier_uniform_(weight)
-
-    def forward(
-        self,
-        x
-    ):
-        for ind, weight in enumerate(self.weights):
-            is_first = ind == 0
-
-            if not is_first:
-                x = F.silu(x)
-
-            x = x @ weight
-
-        return x
-
-# memory mlp, but with gated residual + final projection
-
-class GatedResidualMemoryMLP(Module):
-    def __init__(
-        self,
-        dim,
-        depth,
-        expansion_factor = 2.
-    ):
-        super().__init__()
-        dim_hidden = int(dim * expansion_factor)
-
-        self.weights = ParameterList([
-            ParameterList([
-                Parameter(torch.randn(dim, dim_hidden)),
-                Parameter(torch.randn(dim_hidden, dim)),
-                Parameter(torch.randn(dim * 2, dim)),
-            ]) for _ in range(depth)
-        ])
-
-        self.final_proj = Parameter(torch.randn(dim, dim))
-
-        for param in self.parameters():
-            nn.init.xavier_uniform_(param)
-
-    def forward(
-        self,
-        x
-    ):
-        for weight1, weight2, to_gates in self.weights:
-            res = x
-
-            hidden = x @ weight1
-            hidden = F.silu(hidden)
-            branch_out = hidden @ weight2
-
-            # gated residual
-
-            gates = cat((branch_out, res), dim = -1) @ to_gates
-            x = res.lerp(branch_out, gates.sigmoid())
-
-        return x @ self.final_proj
-
-# memory mlp with factorized weights
-# so can tradeoff capacity for smaller chunk sizes
-
-class FactorizedMemoryMLP(Module):
-    def __init__(
-        self,
-        dim,
-        depth,
-        k = 32
-    ):
-        super().__init__()
-        self.weights = ParameterList([
-            ParameterList([
-                Parameter(torch.randn(dim, k)),
-                Parameter(torch.randn(k, dim)),
-            ]) for _ in range(depth)
-        ])
-
-        for weight1, weight2 in self.weights:
-            nn.init.xavier_uniform_(weight1)
-            nn.init.xavier_uniform_(weight2)
-
-    def forward(
-        self,
-        x
-    ):
-        for ind, (weight1, weight2) in enumerate(self.weights):
-            is_first = ind == 0
-
-            if not is_first:
-                x = F.silu(x)
-
-            x = x @ weight1 @ weight2
-
-        return x
-
-# improvised attention as memory module
-
-class MemoryAttention(Module):
-    def __init__(
-        self,
-        dim,
-        scale = 8.,
-        expansion_factor = 2.
-    ):
-        super().__init__()
-        self.scale = scale
-        dim_ff_hidden = int(dim * expansion_factor)
-
-        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
-            nn.Parameter(torch.randn(dim, dim_ff_hidden)), # ff w1
-            nn.Parameter(torch.randn(dim_ff_hidden, dim)), # ff w2
-        ])
-
-        for weight in self.weights:
-            nn.init.xavier_uniform_(weight)
-
-    def forward(self, x):
-        wq, wk, wv, ffw1, ffw2 = self.weights
-
-        q = F.normalize(x @ wq, dim = -1)
-        k = F.normalize(x @ wk, dim = -1)
-        v = x @ wv
-
-        attn_out = F.scaled_dot_product_attention(
-            q, k, v,
-            scale = self.scale,
-            is_causal = True
-        )
-
-        x = x + attn_out
-
-        h = F.silu(x @ ffw1)
-        out = h @ ffw2
-
-        return out
-
 # associative scan wrapper
 
 class AssocScan(Module):
@@ -870,6 +723,8 @@ class NeuralMemory(Module):
             prev_layer_updates = TensorDict(prev_layer_updates)
             prev_layer_updates = prev_layer_updates.apply(lambda t: t[:, -1:])
 
+        values = None
+
         if store_seq_cache_len == self.chunk_size:
 
             next_updates, next_states, values = self.store_memories(
@@ -917,10 +772,13 @@ class NeuralMemory(Module):
         if seq_len < self.retrieve_chunk_size:
             out = self.init_empty_memory_embed(batch, seq_len)
 
-            next_store_state = (seq_len, seq, None, None)
+            next_store_state = NeuralMemCache(seq_len, seq, None, None)
 
             out = (out, next_store_state)
 
+            if return_values:
+                out = (*out, self.zero)
+
             if not return_aux_kv_loss:
                 return out
 

{titans_pytorch-0.1.35.dist-info → titans_pytorch-0.1.37.dist-info}/METADATA

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

titans_pytorch-0.1.37.dist-info/RECORD

@@ -0,0 +1,9 @@
+titans_pytorch/__init__.py,sha256=Y3m_ZlpEqYwp-Md1ARhNGJxq8bQp8ty1o039nZOOJo0,276
+titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
+titans_pytorch/mac_transformer.py,sha256=KBwo-Fr_fDzVaAa7xg1ggEpNlE4vRUoGMEjB-I2ZWTU,26463
+titans_pytorch/memory_models.py,sha256=LI9T36XB6YXIvvGWRw0ZMDlGpRC6KIv03OPzME2VAaU,3772
+titans_pytorch/neural_memory.py,sha256=B6nTdsq7Tp6lfmUwQakOtODImwZzMCDkDfIv5CIKlbQ,25453
+titans_pytorch-0.1.37.dist-info/METADATA,sha256=BdkDj71kq320M_vbZOGPrRBR91ycgqBJaUNm9m7tA5U,6826
+titans_pytorch-0.1.37.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.1.37.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.1.37.dist-info/RECORD,,

titans_pytorch-0.1.35.dist-info/RECORD

@@ -1,8 +0,0 @@
-titans_pytorch/__init__.py,sha256=rMT99CPQFH4Gudp0FmVPWGKfhBf6xksGEaFEcOVdqjs,230
-titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
-titans_pytorch/mac_transformer.py,sha256=KBwo-Fr_fDzVaAa7xg1ggEpNlE4vRUoGMEjB-I2ZWTU,26463
-titans_pytorch/neural_memory.py,sha256=wFOLFe3ViXiQfNvUiAGJ6BfiaDhr0BYDRDnLNMHWQhU,28938
-titans_pytorch-0.1.35.dist-info/METADATA,sha256=5e5qPt4hAOhxhDWqdjutjJuUmht44zYq_KqgagKjqxE,6826
-titans_pytorch-0.1.35.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.1.35.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.1.35.dist-info/RECORD,,