titans-pytorch 0.0.40__py3-none-any.whl → 0.0.41__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (5) hide show
  1. titans_pytorch/mac_transformer.py +4 -1
  2. {titans_pytorch-0.0.40.dist-info → titans_pytorch-0.0.41.dist-info}/METADATA +1 -1
  3. {titans_pytorch-0.0.40.dist-info → titans_pytorch-0.0.41.dist-info}/RECORD +5 -5
  4. {titans_pytorch-0.0.40.dist-info → titans_pytorch-0.0.41.dist-info}/WHEEL +0 -0
  5. {titans_pytorch-0.0.40.dist-info → titans_pytorch-0.0.41.dist-info}/licenses/LICENSE +0 -0
titans_pytorch/mac_transformer.py CHANGED
@@ -192,6 +192,7 @@ class MemoryAsContextTransformer(Module):
192
192
  dim,
193
193
  depth,
194
194
  segment_len,
195
+ neural_memory_segment_len = None,
195
196
  num_longterm_mem_tokens = 0,
196
197
  num_persist_mem_tokens = 0,
197
198
  dim_head = 64,
@@ -221,7 +222,9 @@ class MemoryAsContextTransformer(Module):
221
222
  init_hyper_conn, self.expand_streams, self.reduce_streams = get_init_and_expand_reduce_stream_functions(num_residual_streams, disable = num_residual_streams == 1)
222
223
 
223
224
  self.layers = ModuleList([])
225
+
224
226
  self.neural_mem_layers = ModuleList([])
227
+ neural_memory_segment_len = default(neural_memory_segment_len, num_longterm_mem_tokens + segment_len)
225
228
 
226
229
  layers = tuple(range(1, depth + 1))
227
230
 
@@ -242,7 +245,7 @@ class MemoryAsContextTransformer(Module):
242
245
 
243
246
  mem = NeuralMemory(
244
247
  dim = dim,
245
- chunk_size = num_longterm_mem_tokens + segment_len,
248
+ chunk_size = neural_memory_segment_len,
246
249
  **neural_memory_kwargs
247
250
  )
248
251
 
{titans_pytorch-0.0.40.dist-info → titans_pytorch-0.0.41.dist-info}/METADATA RENAMED
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: titans-pytorch
3
- Version: 0.0.40
3
+ Version: 0.0.41
4
4
  Summary: Titans
5
5
  Project-URL: Homepage, https://pypi.org/project/titans-pytorch/
6
6
  Project-URL: Repository, https://github.com/lucidrains/titans-pytorch
{titans_pytorch-0.0.40.dist-info → titans_pytorch-0.0.41.dist-info}/RECORD RENAMED
@@ -1,9 +1,9 @@
1
1
  titans_pytorch/__init__.py,sha256=F6pV8BamKCsbJFVo5x2hw69vzfJNLy54SwIKIueMdp4,142
2
2
  titans_pytorch/associative_scan.py,sha256=Y-iYqmFuG-NoCKu6kgql1mhowXTeJfyawi3eUIXamp0,2650
3
- titans_pytorch/mac_transformer.py,sha256=y9sruSvGCEL4flu_RW7bCdvIe-S9dEdGacbmPYL1kqA,9311
3
+ titans_pytorch/mac_transformer.py,sha256=szHg8m97ew7OtipVih3pkOe1jsvhBnqvohJVJBrU5ks,9452
4
4
  titans_pytorch/titans.py,sha256=bv2Ceq-_4nNb5FNx4hLd2inC93m5MmJxO2-Mbf6PKK4,14378
5
5
  titans_pytorch/titans_attn_memory.py,sha256=Rwx_-riGeISBefZg5Kjic8jzmmRRys-u93D2Kgb7Mos,12691
6
- titans_pytorch-0.0.40.dist-info/METADATA,sha256=JCJ5aG9_-rVUErW6u-DXkJtVQ52Bf3XQDN3puirXAXo,3968
7
- titans_pytorch-0.0.40.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
8
- titans_pytorch-0.0.40.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
9
- titans_pytorch-0.0.40.dist-info/RECORD,,
6
+ titans_pytorch-0.0.41.dist-info/METADATA,sha256=RXjWshmGlzoA-OYaiI6GwieXBwbsq5ZR_hk-PM66aeg,3968
7
+ titans_pytorch-0.0.41.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
8
+ titans_pytorch-0.0.41.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
9
+ titans_pytorch-0.0.41.dist-info/RECORD,,