misunderstood how activation functions were applied

Author: lucidrains

gateloop_transformer/gateloop_transformer.py

@@ -139,15 +139,16 @@ def binary_operator(a, b):
         a_i, kv_i = a
         a_j, kv_j = b
 
-        # unsure, but i think this is what the paper was doing
-        # feel free to open an issue if not
-
-        a_i = a_i.real.sigmoid() + 1.j * a_i.imag
-        a_j = a_j.real.sigmoid() + 1.j * a_j.imag
-
         return a_j * a_i, a_j.real * kv_i + kv_j
 
     a = rearrange(a, '... -> ... 1')
+
+    # activations for state transitions
+    # sigmoid for magnitude, identity for phase
+
+    magnitude, phase = a.abs(), a.angle()
+    a = torch.polar(magnitude.sigmoid(), phase)
+
     _, kv = associative_scan(binary_operator, (a, kv), axis = 1)
 
     return einsum('b n d, b n d e -> b n e', q, kv)

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'gateloop-transformer',
   packages = find_packages(exclude=[]),
-  version = '0.0.6',
+  version = '0.0.7',
   license='MIT',
   description = 'GateLoop Transformer',
   author = 'Phil Wang',