address #4

Author: lucidrains

gateloop_transformer/simplified_gate_loop.py

@@ -17,17 +17,26 @@
 def exists(v):
     return v is not None
 
-def abs_clamp_eps(t, eps = 1e-20):
+def default(v, d):
+    return v if exists(v) else d
+
+def eps_by_dtype(dtype):
+    return 1e-7 if dtype == torch.float16 else 1e-20
+
+def abs_clamp_eps(t, eps = None):
+    eps = default(eps, eps_by_dtype(t.dtype))
     sign = torch.sign(t)
     return sign * t.abs().clamp(min = eps)
 
 # associative scan using heinsen sequences
 # 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, eps = 1e-20):
+def heinsen_associative_scan(a, kv, eps = None):
+    eps = default(eps, eps_by_dtype(a.dtype))
     log_a = a.clamp(min = eps).log()
-    log_kv = abs_clamp_eps(kv, eps = eps).to(dtype = torch.complex64).log()
+
+    log_kv = abs_clamp_eps(kv).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.4',
+  version = '0.2.5',
   license='MIT',
   description = 'GateLoop Transformer',
   author = 'Phil Wang',