dep torch 2.1+

gateloop_transformer/simplified_gate_loop.py

@@ -25,9 +25,9 @@ def abs_clamp_eps(t, eps = 1e-20):
 # https://github.com/glassroom/heinsen_sequence
 # graciously shared to the world by Franz A. Heinsen in https://arxiv.org/abs/2311.06281 in October 2023
 
-def heinsen_associative_scan(a, kv):
-    log_a = a.clamp(min = 1e-20).log()
-    log_kv = abs_clamp_eps(kv).to(dtype = torch.complex64).log()
+def heinsen_associative_scan(a, kv, eps = 1e-20):
+    log_a = a.clamp(min = eps).log()
+    log_kv = abs_clamp_eps(kv, eps = eps).to(dtype = torch.complex64).log()
 
     a_star = torch.cumsum(log_a, dim = 1)
     log_x0_plus_b_star = torch.logcumsumexp(log_kv - a_star, dim = 1)

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'gateloop-transformer',
   packages = find_packages(exclude=[]),
-  version = '0.2.1',
+  version = '0.2.2',
   license='MIT',
   description = 'GateLoop Transformer',
   author = 'Phil Wang',
@@ -18,7 +18,7 @@
   install_requires=[
     'einops>=0.7.0',
     'rotary-embedding-torch',
-    'torch>=2.0',
+    'torch>=2.1',
   ],
   classifiers=[
     'Development Status :: 4 - Beta',