titans-pytorch 0.1.8__py3-none-any.whl → 0.1.9__py3-none-any.whl

titans_pytorch/mac_transformer.py

@@ -83,6 +83,14 @@ def identity(t):
 def round_up_multiple(seq, mult):
     return ceil(seq / mult) * mult
 
+def pack_with_inverse(t, pattern):
+    packed, packed_shape = pack(t, pattern)
+
+    def inverse(out, inv_pattern = None):
+        return unpack(out, packed_shape, default(inv_pattern, pattern))
+
+    return packed, inverse
+
 def pad_at_dim(t, pad, dim = -1, value = 0.):
     dims_from_right = (- dim - 1) if dim < 0 else (t.ndim - dim - 1)
     zeros = ((0, 0) * dims_from_right)
@@ -576,7 +584,7 @@ class MemoryAsContextTransformer(Module):
         x, inverse_segment = pad_and_segment_with_inverse(x, segment_len)
 
         mems = repeat(self.longterm_mems, 'n d -> b n d', b = x.shape[0])
-        x, mem_ps = pack((x, mems), 'b * d')
+        x, inverse_pack_mems = pack_with_inverse((x, mems), 'b * d')
 
         x = inverse_segment(x)
 
@@ -634,7 +642,7 @@ class MemoryAsContextTransformer(Module):
 
         x, inverse_segment = pad_and_segment_with_inverse(x, segment_len + num_longterm_mem_tokens)
 
-        x, _ = unpack(x, mem_ps, 'b * d')
+        x, _ = inverse_pack_mems(x)
 
         x = inverse_segment(x)
 

titans_pytorch/titans.py

@@ -391,6 +391,8 @@ class NeuralMemory(Module):
 
         # whether to use averaging of chunks, or attention pooling
 
+        assert not (attn_pool_chunks and chunk_size == 1), '`attn_pool_chunks` cannot be set to True if `chunk_size` is set to 1'
+
         if not attn_pool_chunks:
             chunk_reduce_module = Reduce('b (n c) ... -> b n ...', 'mean', c = chunk_size)
         else:

{titans_pytorch-0.1.8.dist-info → titans_pytorch-0.1.9.dist-info}/METADATA

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

titans_pytorch-0.1.9.dist-info/RECORD

@@ -0,0 +1,8 @@
+titans_pytorch/__init__.py,sha256=u0tta_KqhOdfzCEDWT9P4_jejJEK2q1XxhsEzB5MnQU,223
+titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
+titans_pytorch/mac_transformer.py,sha256=hAKuyu-dfSD23hZkCJDPf3PmRSWouVr-XLIDkcDp1MU,19364
+titans_pytorch/titans.py,sha256=gZvYk1j6aBMp0uE6l1a2GH_4ea9W2uXKytJb3CDPTlk,21162
+titans_pytorch-0.1.9.dist-info/METADATA,sha256=J-sURfBOBbNykVWYRPE4RXayhVzEE5vcTStAoznuasM,4747
+titans_pytorch-0.1.9.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.1.9.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.1.9.dist-info/RECORD,,

titans_pytorch-0.1.8.dist-info/RECORD

@@ -1,8 +0,0 @@
-titans_pytorch/__init__.py,sha256=u0tta_KqhOdfzCEDWT9P4_jejJEK2q1XxhsEzB5MnQU,223
-titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
-titans_pytorch/mac_transformer.py,sha256=YowJzQC2p3fWgzFCe9jXrw1b3wWPKN-dcLs5sX-3r8Y,19123
-titans_pytorch/titans.py,sha256=qRUw-Lad_dkMqV7ASMNoGLgxYwGD-maAadetAd_qmc8,21031
-titans_pytorch-0.1.8.dist-info/METADATA,sha256=0-m6h7GERineU8N9_2cW6nCuXs96twFEwVYkHVuuuLM,4747
-titans_pytorch-0.1.8.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.1.8.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.1.8.dist-info/RECORD,,