[Simple GLA] Add parallel form impls

fla-org · Nov 9, 2024 · 3a5bb44 · 3a5bb44
1 parent 5ff82ac
commit 3a5bb44
Show file tree

Hide file tree

Showing 3 changed files with 625 additions and 3 deletions.
diff --git a/fla/ops/simple_gla/__init__.py b/fla/ops/simple_gla/__init__.py
@@ -2,8 +2,10 @@
 
 from .chunk import chunk_simple_gla
 from .fused_recurrent import fused_recurrent_simple_gla
+from .parallel import ParallelSimpleGLAFunction
 
 __all__ = [
     'chunk_simple_gla',
-    'fused_recurrent_simple_gla'
+    'fused_recurrent_simple_gla',
+    'ParallelSimpleGLAFunction'
 ]