titans-pytorch 0.0.52__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/mac_transformer.py

@@ -1,4 +1,5 @@
 from __future__ import annotations
+from typing import Callable
 from math import ceil
 from functools import partial
 
@@ -32,7 +33,7 @@ def create_mac_block_mask(seq_len, window_size, persist_mem_len):
 
 # einstein notation related
 
-from einops import einsum, repeat, rearrange, pack, unpack
+from einops import repeat, rearrange, pack, unpack
 from einops.layers.torch import Rearrange
 
 # b - batch

titans_pytorch/titans.py

@@ -1,4 +1,5 @@
 from __future__ import annotations
+from typing import Callable
 import math
 from functools import partial
 
@@ -16,7 +17,6 @@ from titans_pytorch.associative_scan import (
     pad_at_dim
 )
 
-import einx
 from einops import rearrange, repeat, pack, unpack
 from einops.layers.torch import Rearrange, Reduce
 
@@ -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
@@ -338,9 +346,9 @@ class NeuralMemory(Module):
 
         # take care of chunking
 
-        keys, values = tuple(rearrange(t, 'b (n c) d -> (b n) c d', c = self.chunk_size) for t in (keys, values))
+        keys, values = tuple(rearrange(t, 'b (n c) d -> (b n) c d', c = chunk_size) for t in (keys, values))
 
-        adaptive_lr = rearrange(adaptive_lr, 'b (n c) -> (b n) c', c = self.chunk_size)
+        adaptive_lr = rearrange(adaptive_lr, 'b (n c) -> (b n) c', c = chunk_size)
 
         # get grads and extra auxiliary loss (for backwarding through qkv projection in base neural memory module)
 

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.0.52
+Version: 0.0.54
 Summary: Titans
 Project-URL: Homepage, https://pypi.org/project/titans-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/titans-pytorch
@@ -37,7 +37,6 @@ Requires-Python: >=3.9
 Requires-Dist: accelerated-scan>=0.2.0
 Requires-Dist: axial-positional-embedding>=0.3.5
 Requires-Dist: einops>=0.8.0
-Requires-Dist: einx>=0.3.0
 Requires-Dist: hyper-connections>=0.1.8
 Requires-Dist: ninja
 Requires-Dist: rotary-embedding-torch

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.52.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=khfjpbsy-uT9NIG3dZLsLOG_XSEi7EqcyfbPr7EQc2Q,13192
-titans_pytorch/titans.py,sha256=T04onF0xhcrosS-Qkx7fcx-Cqgh0TdU5JLdq9l8ayGg,15911
-titans_pytorch-0.0.52.dist-info/METADATA,sha256=coC9ExIuNvmab0BktSE1NwUgxRaBUV7h_cTHeoJkRJo,4484
-titans_pytorch-0.0.52.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.0.52.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.0.52.dist-info/RECORD,,