until torch supports assoc scan

lucidrains · lucidrains · commit 9f12637d376b · 2025-04-13T06:34:02.000-07:00
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "titans-pytorch"
-version = "0.4.6"
+version = "0.4.7"
 description = "Titans"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -25,7 +25,7 @@ classifiers=[
 ]
 
 dependencies = [
-    "accelerated-scan>=0.2.0",
+    "assoc-scan",
     "axial_positional_embedding>=0.3.10",
     "einops>=0.8.0",
     "einx>=0.3.0",
diff --git a/titans_pytorch/associative_scan.py b/titans_pytorch/associative_scan.py
diff --git a/titans_pytorch/neural_memory.py b/titans_pytorch/neural_memory.py
@@ -15,7 +15,7 @@
 
 from tensordict import TensorDict
 
-from titans_pytorch.associative_scan import AssocScan
+from assoc_scan import AssocScan
 
 from titans_pytorch.memory_models import(
     MemoryMLP,
