titans-pytorch 0.4.1__py3-none-any.whl → 0.4.3__py3-none-any.whl

titans_pytorch/neural_memory.py

@@ -652,6 +652,7 @@ class NeuralMemory(Module):
 
         # surprises
 
+        adaptive_lr = rearrange(adaptive_lr, '(b h n) c -> b h (n c)', b = batch, h = heads)
         unweighted_mem_model_loss = rearrange(unweighted_mem_model_loss, '(b h n) c -> b h (n c)', b = batch, h = heads)
 
         # maybe softclamp grad norm
@@ -695,7 +696,7 @@ class NeuralMemory(Module):
             if not return_surprises:
                 return output
 
-            return (*output, unweighted_mem_model_loss)
+            return (*output, (unweighted_mem_model_loss, adaptive_lr))
 
         # momentum + weight decay - momentum is the new contribution, as most linear RNNs have learned forgetting gates
 
@@ -755,7 +756,7 @@ class NeuralMemory(Module):
         if not return_surprises:
             return updates, next_store_state
 
-        return updates, next_store_state, unweighted_mem_model_loss
+        return updates, next_store_state, (unweighted_mem_model_loss, adaptive_lr)
 
     def retrieve_memories(
         self,
@@ -939,7 +940,7 @@ class NeuralMemory(Module):
 
         # whether to allow network to slowly adjust from initial weight throughout (residual path) to fully updating weights every batch
 
-        surprises = None
+        surprises = (None, None)
         gate = None
 
         if exists(self.transition_gate):
@@ -966,7 +967,7 @@ class NeuralMemory(Module):
 
             updates = accum_updates(updates, next_updates)
 
-            surprises = safe_cat((surprises, chunk_surprises), dim = -1)
+            surprises = tuple(safe_cat(args, dim = -1) for args in zip(surprises, chunk_surprises))
 
             if is_last and not update_after_final_store:
                 continue

{titans_pytorch-0.4.1.dist-info → titans_pytorch-0.4.3.dist-info}/METADATA

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

{titans_pytorch-0.4.1.dist-info → titans_pytorch-0.4.3.dist-info}/RECORD

@@ -2,8 +2,8 @@ titans_pytorch/__init__.py,sha256=AyEUlcXWpnqrvyeihRAXWIfQlzLA4NhBjOqQU4edL-4,29
 titans_pytorch/associative_scan.py,sha256=esaLbukFlgvy2aqopsqBy6KEcZ64B3rsNhG8moKdPSc,5159
 titans_pytorch/mac_transformer.py,sha256=tz72141G5t3AOnxSVsOLtLptGtl8T7zROUvaTw2_XCY,26960
 titans_pytorch/memory_models.py,sha256=wnH9i9kUSoVZhEWUlj8LpBSbB400L9kLt1zP8CO45QQ,5835
-titans_pytorch/neural_memory.py,sha256=io5fvLWpOTzx8mkDA9sg3Mkc7-aeugUJoDCniryiuYE,32666
-titans_pytorch-0.4.1.dist-info/METADATA,sha256=XwduHOXOJvjaWJhdYUq-1jhVq2zNKJBwMH1VWopxv5Y,6816
-titans_pytorch-0.4.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-titans_pytorch-0.4.1.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-titans_pytorch-0.4.1.dist-info/RECORD,,
+titans_pytorch/neural_memory.py,sha256=HdBaRGURJ84Qy-a6PdfeQoc5ZzY7H0c5YHUASaSVu1A,32824
+titans_pytorch-0.4.3.dist-info/METADATA,sha256=SIq5KS2xehsUAwuFpRSFNdnLbgamWUMLN5xj4MJGRe0,6816
+titans_pytorch-0.4.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+titans_pytorch-0.4.3.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+titans_pytorch-0.4.3.dist-info/RECORD,,