titans-pytorch 0.0.17__py3-none-any.whl → 0.0.19__py3-none-any.whl

titans_pytorch/titans.py

@@ -17,7 +17,7 @@ from titans_pytorch.associative_scan import (
 )
 
 import einx
-from einops import rearrange, pack, unpack
+from einops import rearrange, repeat, pack, unpack
 from einops.layers.torch import Rearrange, Reduce
 
 """
@@ -152,6 +152,11 @@ class NeuralMemory(Module):
         self.to_keys_values = LinearNoBias(dim, dim_inner * 2)
         self.store_memory_loss_fn = store_memory_loss_fn
 
+        # empty memory embed
+
+        self.empty_memory_embed = nn.Parameter(torch.zeros(dim))
+        nn.init.normal_(self.empty_memory_embed, std = 0.02)
+
         # learned adaptive learning rate and momentum
         # todo - explore mlp layerwise learned lr / momentum
 
@@ -187,6 +192,9 @@ class NeuralMemory(Module):
 
         return init_weights, init_momentum
 
+    def init_empty_memory_embed(self, batch, seq_len):
+        return repeat(self.empty_memory_embed, 'd -> b n d', b = batch, n = seq_len)
+
     def store_memories(
         self,
         seq,
@@ -269,7 +277,7 @@ class NeuralMemory(Module):
                 gates = F.pad(gates, (0, next_power_two_seq_len - seq_len))
                 inputs = F.pad(inputs, (0, next_power_two_seq_len - seq_len))
 
-                outputs = scan(gates, inputs)
+                outputs = scan(gates.contiguous(), inputs.contiguous())
 
                 outputs = outputs[..., :seq_len]
                 outputs = rearrange(outputs, 'b d n -> b n d')
@@ -372,11 +380,12 @@ class NeuralMemory(Module):
 
         values = self.post_rmsnorm(values)
 
-        # restore
+        # restore, pad with empty memory embed
 
-        values = pad_at_dim(values, (chunk_size - 1, 0), dim = 1, value = 0.) # todo, used a learned null memory embedding instead of 0s for retrieving from empty neural memory
-        values = values[:, :-padding]
+        empty_memory_embeds = self.init_empty_memory_embed(values.shape[0], chunk_size - 1)
+        values = torch.cat((empty_memory_embeds, values), dim = -2)
 
+        values = values[:, :-padding]
         return values
 
     def forward(
@@ -389,7 +398,7 @@ class NeuralMemory(Module):
         batch, seq_len = seq.shape[:2]
 
         if seq_len < self.chunk_size:
-            return torch.zeros_like(seq)
+            return self.init_empty_memory_embed(batch, seq_len)
 
         if exists(past_state):
             past_state = tuple(TensorDict(d) for d in past_state)

{titans_pytorch-0.0.17.dist-info → titans_pytorch-0.0.19.dist-info}/METADATA

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

titans_pytorch-0.0.19.dist-info/RECORD

@@ -0,0 +1,8 @@
+titans_pytorch/__init__.py,sha256=nB0873FZ_OyCda3qFeWTdpO4LbbrXDEvtAefVLzh6g0,71
+titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
+titans_pytorch/titans.py,sha256=CxbJgNdIS9NbbCDdgotFXAnrV16xmvufUErerKe7qJA,12636
+titans_pytorch/titans_attn_memory.py,sha256=Rwx_-riGeISBefZg5Kjic8jzmmRRys-u93D2Kgb7Mos,12691
+titans_pytorch-0.0.19.dist-info/METADATA,sha256=5Wpk79HYI4z8LeNRV__UaamKppiGcJ2HdIlll1JSZr8,3811
+titans_pytorch-0.0.19.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.0.19.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.0.19.dist-info/RECORD,,

titans_pytorch-0.0.17.dist-info/RECORD

@@ -1,8 +0,0 @@
-titans_pytorch/__init__.py,sha256=nB0873FZ_OyCda3qFeWTdpO4LbbrXDEvtAefVLzh6g0,71
-titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
-titans_pytorch/titans.py,sha256=HYm0R_1w3s8MNPsyE2qAVpHGqTBX_AoWtjzxRfF1Ams,12269
-titans_pytorch/titans_attn_memory.py,sha256=Rwx_-riGeISBefZg5Kjic8jzmmRRys-u93D2Kgb7Mos,12691
-titans_pytorch-0.0.17.dist-info/METADATA,sha256=s8PEQdaW8WSjZkjlho1Gv1gcsk7GD2lu9oka7bP5Rf8,3811
-titans_pytorch-0.0.17.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.0.17.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.0.17.dist-info/RECORD,,