titans-pytorch 0.0.49__tar.gz → 0.0.50__tar.gz

{titans_pytorch-0.0.49 → titans_pytorch-0.0.50}/PKG-INFO

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

{titans_pytorch-0.0.49 → titans_pytorch-0.0.50}/pyproject.toml

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

{titans_pytorch-0.0.49 → titans_pytorch-0.0.50}/titans_pytorch/titans.py

@@ -162,7 +162,8 @@ class NeuralMemory(Module):
         heads = 1,
         model: Module | None = None,
         store_memory_loss_fn: Callable = default_loss_fn,
-        adaptive_step_transform: Callable = default_adaptive_step_transform,
+        adaptive_step_transform: Callable | None = None,
+        default_step_transform_max_lr = 1e-2,
         pre_rmsnorm = True,
         post_rmsnorm = True,
         max_grad_norm: float | None = None,
@@ -250,6 +251,9 @@ class NeuralMemory(Module):
             Rearrange('b n h -> (b h) n')
         )
 
+        if not exists(adaptive_step_transform):
+            adaptive_step_transform = partial(default_adaptive_step_transform, max_lr = default_step_transform_max_lr)
+
         self.adaptive_step_transform = adaptive_step_transform
 
         # allow for softclamp the gradient norms for storing memories

{titans_pytorch-0.0.49 → titans_pytorch-0.0.50}/train_mac.py

@@ -30,8 +30,8 @@ NUM_PERSIST_MEM = 4
 NUM_LONGTERM_MEM = 4
 NEURAL_MEM_LAYERS = (2, 4)
 WINDOW_SIZE = 32
-KV_RECON_LOSS_WEIGHT = 0.1
-RUN_NAME = f'mac - {NUM_LONGTERM_MEM} longterm mems, layers {NEURAL_MEM_LAYERS} - kv recon loss 0.1'
+KV_RECON_LOSS_WEIGHT = 0.
+RUN_NAME = f'mac - {NUM_LONGTERM_MEM} longterm mems, layers {NEURAL_MEM_LAYERS}'
 
 # wandb experiment tracker