titans-pytorch 0.4.3__py3-none-any.whl → 0.4.4__py3-none-any.whl

titans_pytorch/__init__.py

@@ -1,5 +1,7 @@
 from titans_pytorch.neural_memory import (
     NeuralMemory,
+    NeuralMemState,
+    mem_state_detach
 )
 
 from titans_pytorch.memory_models import (

titans_pytorch/neural_memory.py

@@ -7,10 +7,11 @@ from itertools import zip_longest
 from collections import namedtuple
 
 import torch
-from torch import nn, stack, cat, tensor, Tensor
+from torch import nn, stack, cat, is_tensor, tensor, Tensor
 import torch.nn.functional as F
 from torch.nn import Linear, Module, Parameter, ParameterList, ParameterDict
 from torch.func import functional_call, vmap, grad
+from torch.utils._pytree import tree_map, tree_flatten, tree_unflatten
 
 from tensordict import TensorDict
 
@@ -40,6 +41,8 @@ u - key / value updates - allowing a token to emit multiple key / values
 
 LinearNoBias = partial(Linear, bias = False)
 
+# neural mem state related
+
 NeuralMemState = namedtuple('NeuralMemState', [
     'seq_index',
     'weights',
@@ -48,6 +51,13 @@ NeuralMemState = namedtuple('NeuralMemState', [
     'updates',
 ])
 
+def mem_state_detach(
+    state: NeuralMemState
+):
+    assert isinstance(state, NeuralMemState)
+    state = tree_map(lambda t: t.detach() if is_tensor(t) else t, tuple(state))
+    return NeuralMemState(*state)
+
 # functions
 
 def exists(v):

{titans_pytorch-0.4.3.dist-info → titans_pytorch-0.4.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.4.3
+Version: 0.4.4
 Summary: Titans
 Project-URL: Homepage, https://pypi.org/project/titans-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/titans-pytorch
@@ -56,7 +56,7 @@ Description-Content-Type: text/markdown
 
 <img src="./fig1.png" width="400px"></img>
 
-## Titans - Pytorch (wip)
+## Titans - Pytorch
 
 Unofficial implementation of [Titans](https://arxiv.org/abs/2501.00663) in Pytorch. Will also contain some explorations into architectures beyond their simple 1-4 layer MLP for the neural memory module, if it works well to any degree.
 

titans_pytorch-0.4.4.dist-info/RECORD

@@ -0,0 +1,9 @@
+titans_pytorch/__init__.py,sha256=sVTOuRUkaIYabFExdLY6s1qXm1UwHHz_J19H8ZV-X74,338
+titans_pytorch/associative_scan.py,sha256=esaLbukFlgvy2aqopsqBy6KEcZ64B3rsNhG8moKdPSc,5159
+titans_pytorch/mac_transformer.py,sha256=tz72141G5t3AOnxSVsOLtLptGtl8T7zROUvaTw2_XCY,26960
+titans_pytorch/memory_models.py,sha256=wnH9i9kUSoVZhEWUlj8LpBSbB400L9kLt1zP8CO45QQ,5835
+titans_pytorch/neural_memory.py,sha256=eDC8rl8S241zp8pxKrIEi_6nFm6CEoaH9K4hnDfgzu8,33145
+titans_pytorch-0.4.4.dist-info/METADATA,sha256=CWciTl1VeOvwyL_lqr0JsdmDDIjXG85N8ykwd3w2TxQ,6810
+titans_pytorch-0.4.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.4.4.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.4.4.dist-info/RECORD,,

titans_pytorch-0.4.3.dist-info/RECORD

@@ -1,9 +0,0 @@
-titans_pytorch/__init__.py,sha256=AyEUlcXWpnqrvyeihRAXWIfQlzLA4NhBjOqQU4edL-4,297
-titans_pytorch/associative_scan.py,sha256=esaLbukFlgvy2aqopsqBy6KEcZ64B3rsNhG8moKdPSc,5159
-titans_pytorch/mac_transformer.py,sha256=tz72141G5t3AOnxSVsOLtLptGtl8T7zROUvaTw2_XCY,26960
-titans_pytorch/memory_models.py,sha256=wnH9i9kUSoVZhEWUlj8LpBSbB400L9kLt1zP8CO45QQ,5835
-titans_pytorch/neural_memory.py,sha256=HdBaRGURJ84Qy-a6PdfeQoc5ZzY7H0c5YHUASaSVu1A,32824
-titans_pytorch-0.4.3.dist-info/METADATA,sha256=SIq5KS2xehsUAwuFpRSFNdnLbgamWUMLN5xj4MJGRe0,6816
-titans_pytorch-0.4.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.4.3.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.4.3.dist-info/RECORD,,