titans-pytorch 0.0.21__tar.gz → 0.0.22__tar.gz

{titans_pytorch-0.0.21 → titans_pytorch-0.0.22}/PKG-INFO

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

{titans_pytorch-0.0.21 → titans_pytorch-0.0.22}/pyproject.toml

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

{titans_pytorch-0.0.21 → titans_pytorch-0.0.22}/tests/test_titans.py

@@ -2,8 +2,10 @@ import torch
 import pytest
 
 @pytest.mark.parametrize('seq_len', (32, 1024, 77))
+@pytest.mark.parametrize('max_grad_norm', (None, 2.))
 def test_titans(
-    seq_len
+    seq_len,
+    max_grad_norm
 ):
 
     from titans_pytorch import NeuralMemory
@@ -11,6 +13,7 @@ def test_titans(
     mem = NeuralMemory(
         dim = 384,
         chunk_size = 64,
+        max_grad_norm = max_grad_norm
     )
 
     seq = torch.randn(2, seq_len, 384)

{titans_pytorch-0.0.21 → titans_pytorch-0.0.22}/titans_pytorch/titans.py

@@ -100,8 +100,8 @@ class MemoryMLP(Module):
 
 # main neural memory
 
-def default_adaptive_step_transform(adaptive_step):
-    return torch.exp(adaptive_step.sigmoid() * -15) # from 1. - 1e-7
+def default_adaptive_step_transform(adaptive_step, max_lr = 1e-1):
+    return adaptive_step.sigmoid() * max_lr
 
 def default_loss_fn(pred, target):
     return (pred - target).pow(2).mean(dim = -1).sum()

{titans_pytorch-0.0.21 → titans_pytorch-0.0.22}/train.py

@@ -67,6 +67,7 @@ titans_neural_memory = NeuralMemory(
     post_rmsnorm = True,
     dim_head = 64,
     heads = 4,
+    max_grad_norm = 1.,
     use_accelerated_scan = True,
     default_mlp_kwargs = dict(
         depth = NEURAL_MEMORY_DEPTH