titans-pytorch 0.0.53__py3-none-any.whl → 0.0.54__py3-none-any.whl

titans_pytorch/__init__.py

@@ -1,6 +1,7 @@
 from titans_pytorch.titans import (
     NeuralMemory,
     MemoryMLP,
+    MemoryAttention
 )
 
 from titans_pytorch.mac_transformer import (

titans_pytorch/titans.py

@@ -123,9 +123,12 @@ class MemoryMLP(Module):
 class MemoryAttention(Module):
     def __init__(
         self,
-        dim
+        dim,
+        scale = 8.
     ):
         super().__init__()
+        self.scale = scale
+
         self.weights = nn.ParameterList([
             nn.Parameter(torch.randn(dim, dim)), # queries
             nn.Parameter(torch.randn(dim, dim)), # keys
@@ -143,6 +146,7 @@ class MemoryAttention(Module):
 
         attn_out = F.scaled_dot_product_attention(
             q, k, v,
+            scale = self.scale,
             is_causal = True
         )
 
@@ -174,6 +178,7 @@ class NeuralMemory(Module):
         default_step_transform_max_lr = 1e-2,
         pre_rmsnorm = True,
         post_rmsnorm = True,
+        learned_mem_model_weights = True,
         max_grad_norm: float | None = None,
         use_accelerated_scan = False,
         activation: Module | None = None,
@@ -212,6 +217,9 @@ class NeuralMemory(Module):
         if not exists(model):
             model = MemoryMLP(dim_head, **default_model_kwargs)
 
+        if not learned_mem_model_weights:
+            model.requires_grad_(False)
+
         assert not exists(next(model.buffers(), None)), 'model cannot have buffers for now'
 
         # the memory is the weights of the model

{titans_pytorch-0.0.53.dist-info → titans_pytorch-0.0.54.dist-info}/METADATA

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

titans_pytorch-0.0.54.dist-info/RECORD

@@ -0,0 +1,8 @@
+titans_pytorch/__init__.py,sha256=Q0MQA3RS8vqzs-KzSGZkvLR7afQ6ZW9uMOq1MeNuFoY,170
+titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
+titans_pytorch/mac_transformer.py,sha256=GdUAYq6MDRGY0l2ESBH_kM01AEzVztiKmWfblSKxBEM,13212
+titans_pytorch/titans.py,sha256=ZKm-LnVKh1Cxs2tSxr4CcY37KroOOmYtTFM2F3Zb8Xg,16122
+titans_pytorch-0.0.54.dist-info/METADATA,sha256=bxbC3NBO4Sjii7DpFPcmNsO9M1kX76vj947H2DeUceg,4457
+titans_pytorch-0.0.54.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.0.54.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.0.54.dist-info/RECORD,,

titans_pytorch-0.0.53.dist-info/RECORD

@@ -1,8 +0,0 @@
-titans_pytorch/__init__.py,sha256=wnv_Cjdjqh_h5IqLkQ8xrTtA2K663ITEn-1JeeHofTo,150
-titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
-titans_pytorch/mac_transformer.py,sha256=GdUAYq6MDRGY0l2ESBH_kM01AEzVztiKmWfblSKxBEM,13212
-titans_pytorch/titans.py,sha256=ALJHJOkxgMZr-zFOZQB8Y7KJr-6FcSw8Jefyp4ElXho,15917
-titans_pytorch-0.0.53.dist-info/METADATA,sha256=o7Shd67cH6O8YWJQcTcTwuDPPHwiVRY2hNIPxu2DEHA,4457
-titans_pytorch-0.0.53.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.0.53.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.0.53.dist-info/RECORD,,