add initial version of event CSR and CSC

chaobrain · Dec 16, 2024 · 36c1851 · 36c1851
1 parent 2bb42d4
commit 36c1851
Show file tree

Hide file tree

Showing 7 changed files with 689 additions and 23 deletions.
diff --git a/brainstate/event/__init__.py b/brainstate/event/__init__.py
@@ -14,14 +14,16 @@
 # ==============================================================================
 
 
+from ._csr import *
 from ._csr_mv import *
-from ._csr_mv import __all__ as __all_csr
 from ._fixedprob_mv import *
-from ._fixedprob_mv import __all__ as __all_fixed_probability
 from ._linear_mv import *
 from ._xla_custom_op import *
-from ._xla_custom_op import __all__ as __all_xla_custom_op
-from ._linear_mv import __all__ as __all_linear
 
-__all__ = __all_fixed_probability + __all_linear + __all_csr + __all_xla_custom_op
-del __all_fixed_probability, __all_linear, __all_csr, __all_xla_custom_op
+__all__ = [
+    'CSRLinear',
+    'FixedProb',
+    'XLACustomOp',
+    'CSR',
+    'CSC',
+]