titans-pytorch 0.1.28__tar.gz → 0.1.30__tar.gz

{titans_pytorch-0.1.28 → titans_pytorch-0.1.30}/PKG-INFO

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

{titans_pytorch-0.1.28 → titans_pytorch-0.1.30}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "titans-pytorch"
-version = "0.1.28"
+version = "0.1.30"
 description = "Titans"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{titans_pytorch-0.1.28 → titans_pytorch-0.1.30}/titans_pytorch/titans.py

@@ -409,6 +409,7 @@ class NeuralMemory(Module):
     ):
         super().__init__()
         dim_head = default(dim_head, dim)
+        assert not (heads == 1 and dim_head != dim)
 
         self.retrieve_chunk_size, self.store_chunk_size = pair(chunk_size)
 
@@ -566,7 +567,7 @@ class NeuralMemory(Module):
     ):
         assert xnor(exists(value_residual), exists(self.learned_value_residual))
 
-        seq_len, chunk_size = seq.shape[-2], default(chunk_size, self.store_chunk_size)
+        seq_len, heads, chunk_size = seq.shape[-2], self.heads, default(chunk_size, self.store_chunk_size)
 
         # handle edge case
 
@@ -645,7 +646,7 @@ class NeuralMemory(Module):
 
         # restore batch and sequence dimension
 
-        grads = grads.apply(lambda t: rearrange(t, '(b n) ... -> b n ...', b = batch))
+        grads = grads.apply(lambda t: rearrange(t, '(b n) ... -> b n ...', b = batch * heads))
 
         # maybe per layer modulation
 

{titans_pytorch-0.1.28 → titans_pytorch-0.1.30}/train_mac.py

@@ -37,7 +37,6 @@ NEURAL_MEM_SEGMENT_LEN = WINDOW_SIZE // 2 # set smaller for more granularity for
 SLIDING_WINDOWS = True
 STORE_ATTN_POOL_CHUNKS = True # whether to use attention pooling for chunk derived momentum, per-layer lr mod, decay
 KV_RECON_LOSS_WEIGHT = 0.
-LEARNED_MEM_MODEL_WEIGHTS = True
 
 # experiment related
 
@@ -91,7 +90,6 @@ model = MemoryAsContextTransformer(
         attn_pool_chunks = STORE_ATTN_POOL_CHUNKS,
         momentum = NEURAL_MEM_MOMENTUM,
         use_accelerated_scan = USE_ACCELERATED_SCAN,
-        learned_mem_model_weights = LEARNED_MEM_MODEL_WEIGHTS,
         default_model_kwargs = dict(
             depth = NEURAL_MEMORY_DEPTH,
         )