titans-pytorch 0.2.0__tar.gz → 0.2.1__tar.gz

{titans_pytorch-0.2.0 → titans_pytorch-0.2.1}/PKG-INFO

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

{titans_pytorch-0.2.0 → titans_pytorch-0.2.1}/pyproject.toml

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

{titans_pytorch-0.2.0 → titans_pytorch-0.2.1}/titans_pytorch/mac_transformer.py

@@ -491,7 +491,8 @@ class MemoryAsContextTransformer(Module):
         aux_kv_recon_loss_weight = 0.,
         use_flex_attn = False,
         sliding_window_attn = False,
-        weight_tie_memory_model = False
+        weight_tie_memory_model = False,
+        prev_neural_mem_update_for_weights = None
     ):
         super().__init__()
 
@@ -533,6 +534,7 @@ class MemoryAsContextTransformer(Module):
             assert exists(neural_memory_model), '`neural_memory_model` must be explicitly set'
 
         self.weight_tie_memory_model = weight_tie_memory_model
+        self.prev_neural_mem_update_for_weights = default(prev_neural_mem_update_for_weights, weight_tie_memory_model)
 
         # value residual learning for neural memory
 
@@ -702,7 +704,7 @@ class MemoryAsContextTransformer(Module):
 
         # math
 
-        batch, seq_len, neural_mem_segment_len, segment_len, num_longterm_mem_tokens, attn_window_size, weight_tie_memory_model = *x.shape, self.neural_memory_segment_len, self.segment_len, self.num_longterm_mem_tokens, self.attn_window_size, self.weight_tie_memory_model
+        batch, seq_len, neural_mem_segment_len, segment_len, num_longterm_mem_tokens, attn_window_size, prev_neural_mem_update_for_weights = *x.shape, self.neural_memory_segment_len, self.segment_len, self.num_longterm_mem_tokens, self.attn_window_size, self.prev_neural_mem_update_for_weights
 
         seq_len_with_mem = self.seq_len_with_longterm_mem(seq_len)
 
@@ -814,7 +816,7 @@ class MemoryAsContextTransformer(Module):
                 if self.mem_add_value_residual:
                     mem_value_residual = next_mem_value_residual
 
-                if weight_tie_memory_model:
+                if prev_neural_mem_update_for_weights:
                     neural_memory_updates = next_neural_mem_cache.updates
 
                 if self.gate_attn_output:

{titans_pytorch-0.2.0 → titans_pytorch-0.2.1}/train_mac.py

@@ -164,6 +164,6 @@ for i in tqdm.tqdm(range(NUM_BATCHES), mininterval = 10., desc = 'training'):
         prime = decode_tokens(inp)
         print(f'%s \n\n %s', (prime, '*' * 100))
 
-        sample = model.sample(inp[None, ...], GENERATE_LENGTH, use_cache = True)
+        sample = model.sample(inp[None, ...], GENERATE_LENGTH, use_cache = USE_FAST_INFERENCE)
         output_str = decode_tokens(sample[0])
         print(output_str)