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

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
 
@@ -338,9 +338,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.53.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: titans-pytorch
-Version: 0.0.52
+Version: 0.0.53
 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.53.dist-info/RECORD

@@ -0,0 +1,8 @@
+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,,

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,,