prepare gate loop transformer for experiments

Author: lucidrains

README.md

@@ -6,6 +6,28 @@ Implementation of <a href="https://arxiv.org/abs/2311.01927">GateLoop</a> Transf
 
 Jax version will be done with the <a href="https://github.com/patrick-kidger/equinox">Equinox</a> framework
 
+## Install
+
+```bash
+$ pip install gateloop-transformr
+```
+
+## Usage
+
+```python
+import torch
+from gateloop_transformer import Transformer
+
+model = Transformer(
+    num_tokens = 256,
+    dim = 624,
+    depth = 6
+)
+
+ids = torch.randint(0, 256, (1, 1024))
+logits = model(ids) # (1, 1024, 256)
+```
+
 ## Citations
 
 ```bibtex

gateloop_transformer/gateloop_transformer.py

@@ -197,23 +197,36 @@ def __init__(
         dim_head = 64,
         heads = 8,
         ff_mult = 4,
+        use_gate_looped_attn = True,
+        dim_gate_looped_attn = None,
         data_dependent_rel_pos = False,
-        frac_gradient_data_dependent_rel_pos = 0.5
+        frac_gradient_state_transition = 0.5
     ):
         super().__init__()
 
         self.token_emb = nn.Embedding(num_tokens, dim)
 
         layers = ModuleList([])
+
         for _ in range(depth):
-            layers.append(ModuleList([
-                CausalFullAttention(
+
+            if use_gate_looped_attn:
+                spatial_mixer = GateLoopedAttention(
+                    dim = dim,
+                    dim_inner = dim_gate_looped_attn,
+                    frac_gradient_state_transition = frac_gradient_state_transition
+                )
+            else:
+                spatial_mixer = CausalFullAttention(
                     dim = dim,
                     dim_head = dim_head,
                     heads = heads,
                     data_dependent_rel_pos = data_dependent_rel_pos,
-                    frac_gradient_data_dependent_rel_pos = frac_gradient_data_dependent_rel_pos
-                ),
+                    frac_gradient_data_dependent_rel_pos = frac_gradient_state_transition
+                )
+
+            layers.append(ModuleList([
+                spatial_mixer,
                 FeedForward(
                     dim = dim,
                     mult = ff_mult

setup.py

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