brainstate 0.1.0__py2.py3-none-any.whl → 0.1.0.post20241122__py2.py3-none-any.whl

benchmark/COBA_2005.py

@@ -0,0 +1,125 @@
+# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+#
+# Implementation of the paper:
+#
+# - Brette, R., Rudolph, M., Carnevale, T., Hines, M., Beeman, D., Bower, J. M., et al. (2007),
+#   Simulation of networks of spiking neurons: a review of tools and strategies., J. Comput. Neurosci., 23, 3, 349–98
+#
+# which is based on the balanced network proposed by:
+#
+# - Vogels, T. P. and Abbott, L. F. (2005), Signal propagation and logic gating in networks of integrate-and-fire neurons., J. Neurosci., 25, 46, 10786–95
+#
+import os
+import sys
+
+sys.path.append('../')
+os.environ['XLA_PYTHON_CLIENT_MEM_FRACTION'] = '0.99'
+os.environ['JAX_TRACEBACK_FILTERING'] = 'off'
+
+
+import jax
+import brainunit as u
+import time
+import brainstate as bst
+
+
+class EINet(bst.nn.DynamicsGroup):
+    def __init__(self, scale):
+        super().__init__()
+        self.n_exc = int(3200 * scale)
+        self.n_inh = int(800 * scale)
+        self.num = self.n_exc + self.n_inh
+        self.N = bst.nn.LIFRef(self.num, V_rest=-60. * u.mV, V_th=-50. * u.mV, V_reset=-60. * u.mV,
+                               tau=20. * u.ms, tau_ref=5. * u.ms,
+                               V_initializer=bst.init.Normal(-55., 2., unit=u.mV))
+        self.E = bst.nn.AlignPostProj(
+            comm=bst.event.FixedProb(self.n_exc, self.num, prob=80 / self.num, weight=0.6 * u.mS),
+            syn=bst.nn.Expon.desc(self.num, tau=5. * u.ms),
+            out=bst.nn.COBA.desc(E=0. * u.mV),
+            post=self.N
+        )
+        self.I = bst.nn.AlignPostProj(
+            comm=bst.event.FixedProb(self.n_inh, self.num, prob=80 / self.num, weight=6.7 * u.mS),
+            syn=bst.nn.Expon.desc(self.num, tau=10. * u.ms),
+            out=bst.nn.COBA.desc(E=-80. * u.mV),
+            post=self.N
+        )
+
+    def init_state(self, *args, **kwargs):
+        self.rate = bst.ShortTermState(u.math.zeros(self.num))
+
+    def update(self, t, inp):
+        with bst.environ.context(t=t):
+            spk = self.N.get_spike() != 0.
+            self.E(spk[:self.n_exc])
+            self.I(spk[self.n_exc:])
+            self.N(inp)
+            self.rate.value += self.N.get_spike()
+
+
+@bst.compile.jit(static_argnums=0)
+def run(scale: float):
+    # network
+    net = EINet(scale)
+    bst.nn.init_all_states(net)
+
+    duration = 1e4 * u.ms
+    # simulation
+    with bst.environ.context(dt=0.1 * u.ms):
+        times = u.math.arange(0. * u.ms, duration, bst.environ.get_dt())
+        bst.compile.for_loop(lambda t: net.update(t, 20. * u.mA), times)
+
+    return net.num, net.rate.value.sum() / net.num / duration.to_decimal(u.second)
+
+
+for s in [1, 2, 4, 6, 8, 10, 20, 40, 60, 80, 100]:
+    jax.block_until_ready(run(s))
+
+    t0 = time.time()
+    n, rate = jax.block_until_ready(run(s))
+    t1 = time.time()
+    print(f'scale={s}, size={n}, time = {t1 - t0} s, firing rate = {rate} Hz')
+
+
+# A6000 NVIDIA GPU
+
+# scale=1, size=4000, time = 2.659956455230713 s, firing rate = 50.62445068359375 Hz
+# scale=2, size=8000, time = 2.7318649291992188 s, firing rate = 50.613040924072266 Hz
+# scale=4, size=16000, time = 2.807222604751587 s, firing rate = 50.60573959350586 Hz
+# scale=6, size=24000, time = 3.026782512664795 s, firing rate = 50.60918045043945 Hz
+# scale=8, size=32000, time = 3.1258811950683594 s, firing rate = 50.607574462890625 Hz
+# scale=10, size=40000, time = 3.172346353530884 s, firing rate = 50.60942840576172 Hz
+# scale=20, size=80000, time = 3.751189947128296 s, firing rate = 50.612369537353516 Hz
+# scale=40, size=160000, time = 5.0217814445495605 s, firing rate = 50.617958068847656 Hz
+# scale=60, size=240000, time = 7.002646207809448 s, firing rate = 50.61948776245117 Hz
+# scale=80, size=320000, time = 9.384576320648193 s, firing rate = 50.618499755859375 Hz
+# scale=100, size=400000, time = 11.69654369354248 s, firing rate = 50.61605453491211 Hz
+
+
+# AMD Ryzen 7 7840HS
+
+# scale=1, size=4000, time = 4.436027526855469 s, firing rate = 50.6119270324707 Hz
+# scale=2, size=8000, time = 8.349745273590088 s, firing rate = 50.612266540527344 Hz
+# scale=4, size=16000, time = 16.39163303375244 s, firing rate = 50.61349105834961 Hz
+# scale=6, size=24000, time = 15.725558042526245 s, firing rate = 50.6125602722168 Hz
+# scale=8, size=32000, time = 21.31995177268982 s, firing rate = 50.61244583129883 Hz
+# scale=10, size=40000, time = 27.811061143875122 s, firing rate = 50.61423873901367 Hz
+# scale=20, size=80000, time = 45.54235219955444 s, firing rate = 50.61320877075195 Hz
+# scale=40, size=160000, time = 82.22228026390076 s, firing rate = 50.61309814453125 Hz
+# scale=60, size=240000, time = 125.44037556648254 s, firing rate = 50.613094329833984 Hz
+# scale=80, size=320000, time = 171.20458459854126 s, firing rate = 50.613365173339844 Hz
+# scale=100, size=400000, time = 215.4547393321991 s, firing rate = 50.6129150390625 Hz

benchmark/CUBA_2005.py

@@ -0,0 +1,149 @@
+# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+#
+# Implementation of the paper:
+#
+# - Brette, R., Rudolph, M., Carnevale, T., Hines, M., Beeman, D., Bower, J. M., et al. (2007),
+#   Simulation of networks of spiking neurons: a review of tools and strategies., J. Comput. Neurosci., 23, 3, 349–98
+#
+# which is based on the balanced network proposed by:
+#
+# - Vogels, T. P. and Abbott, L. F. (2005), Signal propagation and logic gating in networks of integrate-and-fire neurons., J. Neurosci., 25, 46, 10786–95
+#
+
+import os
+import sys
+
+sys.path.append('../')
+os.environ['XLA_PYTHON_CLIENT_MEM_FRACTION'] = '0.99'
+os.environ['JAX_TRACEBACK_FILTERING'] = 'off'
+
+
+import jax
+import time
+
+import brainunit as u
+
+import brainstate as bst
+
+
+
+class FixedProb(bst.nn.Module):
+    def __init__(self, n_pre, n_post, prob, weight):
+        super().__init__()
+        self.prob = prob
+        self.weight = weight
+        self.n_pre = n_pre
+        self.n_post = n_post
+
+        self.mask = bst.random.rand(n_pre, n_post) < prob
+
+    def update(self, x):
+        return (x @ self.mask) * self.weight
+
+
+class EINet(bst.nn.DynamicsGroup):
+    def __init__(self, scale=1.0):
+        super().__init__()
+        self.n_exc = int(3200 * scale)
+        self.n_inh = int(800 * scale)
+        self.num = self.n_exc + self.n_inh
+        self.N = bst.nn.LIFRef(
+            self.num, V_rest=-49. * u.mV, V_th=-50. * u.mV, V_reset=-60. * u.mV,
+            tau=20. * u.ms, tau_ref=5. * u.ms,
+            V_initializer=bst.init.Normal(-55., 2., unit=u.mV)
+        )
+        self.E = bst.nn.AlignPostProj(
+            comm=bst.event.FixedProb(self.n_exc, self.num, prob=80 / self.num, weight=1.62 * u.mS),
+            # comm=FixedProb(self.n_exc, self.num, prob=80 / self.num, weight=1.62 * u.mS),
+            syn=bst.nn.Expon.desc(self.num, tau=5. * u.ms),
+            out=bst.nn.CUBA.desc(scale=u.volt),
+            post=self.N
+        )
+        self.I = bst.nn.AlignPostProj(
+            comm=bst.event.FixedProb(self.n_inh, self.num, prob=80 / self.num, weight=-9.0 * u.mS),
+            # comm=FixedProb(self.n_inh, self.num, prob=80 / self.num, weight=-9.0 * u.mS),
+            syn=bst.nn.Expon.desc(self.num, tau=10. * u.ms),
+            out=bst.nn.CUBA.desc(scale=u.volt),
+            post=self.N
+        )
+
+    def init_state(self, *args, **kwargs):
+        self.rate = bst.ShortTermState(u.math.zeros(self.num))
+
+    def update(self, t, inp):
+        with bst.environ.context(t=t):
+            spk = self.N.get_spike()
+            self.E(spk[:self.n_exc])
+            self.I(spk[self.n_exc:])
+            self.N(inp)
+            self.rate.value += self.N.get_spike()
+
+
+@bst.compile.jit(static_argnums=0)
+def run(scale: float):
+    # network
+    net = EINet(scale)
+    bst.nn.init_all_states(net)
+
+    duration = 1e4 * u.ms
+    # simulation
+    with bst.environ.context(dt=0.1 * u.ms):
+        times = u.math.arange(0. * u.ms, duration, bst.environ.get_dt())
+        bst.compile.for_loop(lambda t: net.update(t, 20. * u.mA), times,
+                             # pbar=bst.compile.ProgressBar(100)
+                             )
+
+    return net.num, net.rate.value.sum() / net.num / duration.to_decimal(u.second)
+
+
+for s in [1, 2, 4, 6, 8, 10, 20, 40, 60, 80, 100]:
+    jax.block_until_ready(run(s))
+
+    t0 = time.time()
+    n, rate = jax.block_until_ready(run(s))
+    t1 = time.time()
+    print(f'scale={s}, size={n}, time = {t1 - t0} s, firing rate = {rate} Hz')
+
+
+# A6000 NVIDIA GPU
+
+# scale=1, size=4000, time = 2.6354849338531494 s, firing rate = 24.982027053833008 Hz
+# scale=2, size=8000, time = 2.6781561374664307 s, firing rate = 23.719463348388672 Hz
+# scale=4, size=16000, time = 2.7448785305023193 s, firing rate = 24.592931747436523 Hz
+# scale=6, size=24000, time = 2.8237478733062744 s, firing rate = 24.159996032714844 Hz
+# scale=8, size=32000, time = 2.9344418048858643 s, firing rate = 24.956790924072266 Hz
+# scale=10, size=40000, time = 3.042517900466919 s, firing rate = 23.644424438476562 Hz
+# scale=20, size=80000, time = 3.6727631092071533 s, firing rate = 24.226743698120117 Hz
+# scale=40, size=160000, time = 4.857396602630615 s, firing rate = 24.329742431640625 Hz
+# scale=60, size=240000, time = 6.812030792236328 s, firing rate = 24.370006561279297 Hz
+# scale=80, size=320000, time = 9.227966547012329 s, firing rate = 24.41067886352539 Hz
+# scale=100, size=400000, time = 11.405697584152222 s, firing rate = 24.32524871826172 Hz
+
+
+# AMD Ryzen 7 7840HS
+
+# scale=1, size=4000, time = 1.1661601066589355 s, firing rate = 22.438201904296875 Hz
+# scale=2, size=8000, time = 3.3255884647369385 s, firing rate = 23.868364334106445 Hz
+# scale=4, size=16000, time = 6.950139999389648 s, firing rate = 24.21693229675293 Hz
+# scale=6, size=24000, time = 10.011993169784546 s, firing rate = 24.240270614624023 Hz
+# scale=8, size=32000, time = 13.027734518051147 s, firing rate = 24.753198623657227 Hz
+# scale=10, size=40000, time = 16.449942350387573 s, firing rate = 24.7176570892334 Hz
+# scale=20, size=80000, time = 30.754598140716553 s, firing rate = 24.119956970214844 Hz
+# scale=40, size=160000, time = 63.6387836933136 s, firing rate = 24.72784996032715 Hz
+# scale=60, size=240000, time = 78.58532166481018 s, firing rate = 24.402742385864258 Hz
+# scale=80, size=320000, time = 102.4250214099884 s, firing rate = 24.59092140197754 Hz
+# scale=100, size=400000, time = 145.35173273086548 s, firing rate = 24.33751106262207 Hz

brainstate/augment/_autograd.py

@@ -45,7 +45,7 @@ from brainstate.typing import PyTree, Missing
 from brainstate.util import PrettyType, PrettyAttr, PrettyRepr
 
 __all__ = [
-    'vector_grad', 'grad', 'jacrev', 'jacfwd', 'jacobian', 'hessian',
+    'GradientTransform', 'vector_grad', 'grad', 'jacrev', 'jacfwd', 'jacobian', 'hessian',
 ]
 
 A = TypeVar('A')
@@ -159,6 +159,9 @@ def _jacfwd(fun, argnums=0, holomorphic=False, has_aux=False, return_value=False
     return jacfun
 
 
+TransformFn = Callable
+
+
 class GradientTransform(PrettyRepr):
     """
     Automatic Differentiation Transformations for the ``State`` system.
@@ -168,11 +171,11 @@ class GradientTransform(PrettyRepr):
     def __init__(
         self,
         target: Callable,
-        transform: Callable,
-        grad_states: Any,
-        argnums: Optional[Union[int, Sequence[int]]],
-        return_value: bool,
-        has_aux: bool,
+        transform: TransformFn,
+        grad_states: Optional[Union[State, Sequence[State], Dict[str, State]]] = None,
+        argnums: Optional[Union[int, Sequence[int]]] = None,
+        return_value: bool = False,
+        has_aux: bool = False,
         transform_params: Optional[Dict[str, Any]] = None,
     ):
         # gradient variables

brainstate/event/__init__.py

@@ -19,7 +19,9 @@ from ._csr import __all__ as __all_csr
 from ._fixed_probability import *
 from ._fixed_probability import __all__ as __all_fixed_probability
 from ._linear import *
+from ._xla_custom_op import *
+from ._xla_custom_op import __all__ as __all_xla_custom_op
 from ._linear import __all__ as __all_linear
 
-__all__ = __all_fixed_probability + __all_linear + __all_csr
-del __all_fixed_probability, __all_linear, __all_csr
+__all__ = __all_fixed_probability + __all_linear + __all_csr + __all_xla_custom_op
+del __all_fixed_probability, __all_linear, __all_csr, __all_xla_custom_op

brainstate/event/_csr.py

@@ -21,12 +21,11 @@ import jax
 import jax.numpy as jnp
 import numpy as np
 
-from brainstate._state import ParamState, State
+from brainstate._state import ParamState
 from brainstate._utils import set_module_as
 from brainstate.init import param
 from brainstate.nn._module import Module
-from brainstate.typing import ArrayLike
-from ._misc import IntScalar
+from brainstate.typing import ArrayLike, Size
 
 __all__ = [
     'CSRLinear',
@@ -39,12 +38,12 @@ class CSRLinear(Module):
 
     Parameters
     ----------
-    n_pre : int
-        Number of pre-synaptic neurons.
-    n_post : int
-        Number of post-synaptic neurons.
+    in_size : Size
+        Number of pre-synaptic neurons, i.e., input size.
+    out_size : Size
+        Number of post-synaptic neurons, i.e., output size.
     weight : float or callable or jax.Array or brainunit.Quantity
-        Maximum synaptic conductance.
+        Maximum synaptic conductance or a function that returns the maximum synaptic conductance.
     name : str, optional
         Name of the module.
     """
@@ -53,8 +52,8 @@ class CSRLinear(Module):
 
     def __init__(
         self,
-        n_pre: IntScalar,
-        n_post: IntScalar,
+        in_size: Size,
+        out_size: Size,
         indptr: ArrayLike,
         indices: ArrayLike,
         weight: Union[Callable, ArrayLike],
@@ -63,10 +62,11 @@ class CSRLinear(Module):
     ):
         super().__init__(name=name)
 
-        self.in_size = n_pre
-        self.out_size = n_post
-        self.n_pre = n_pre
-        self.n_post = n_post
+        # network size
+        self.in_size = in_size
+        self.out_size = out_size
+        self.n_pre = self.in_size[-1]
+        self.n_post = self.out_size[-1]
 
         # gradient mode
         assert grad_mode in ['vjp', 'jvp'], f"Unsupported grad_mode: {grad_mode}"
@@ -77,9 +77,10 @@ class CSRLinear(Module):
         indices = jnp.asarray(indices)
         assert indptr.ndim == 1, f"indptr must be 1D. Got: {indptr.ndim}"
         assert indices.ndim == 1, f"indices must be 1D. Got: {indices.ndim}"
-        assert indptr.size == n_pre + 1, f"indptr must have size {n_pre + 1}. Got: {indptr.size}"
-        self.indptr = u.math.asarray(indptr)
-        self.indices = u.math.asarray(indices)
+        assert indptr.size == self.n_pre + 1, f"indptr must have size {self.n_pre + 1}. Got: {indptr.size}"
+        with jax.ensure_compile_time_eval():
+            self.indptr = u.math.asarray(indptr)
+            self.indices = u.math.asarray(indices)
 
         # maximum synaptic conductance
         weight = param(weight, (len(indices),), allow_none=False)
@@ -88,7 +89,9 @@ class CSRLinear(Module):
         self.weight = ParamState(weight)
 
     def update(self, spk: jax.Array) -> Union[jax.Array, u.Quantity]:
-        weight = self.weight.value if isinstance(self.weight, State) else self.weight
+        weight = self.weight.value
+
+        # return zero if no pre-synaptic neurons
         if len(self.indices) == 0:
             r = u.math.zeros(spk.shape[:-1] + (self.n_post,),
                              dtype=weight.dtype,
@@ -96,6 +99,8 @@ class CSRLinear(Module):
             return u.maybe_decimal(r)
 
         device_kind = jax.devices()[0].platform  # spk.device.device_kind
+
+        # CPU implementation
         if device_kind == 'cpu':
             return cpu_event_csr(
                 u.math.asarray(spk),
@@ -104,8 +109,11 @@ class CSRLinear(Module):
                 u.math.asarray(weight),
                 n_post=self.n_post, grad_mode=self.grad_mode
             )
+
+        # GPU/TPU implementation
         elif device_kind in ['gpu', 'tpu']:
             raise NotImplementedError()
+
         else:
             raise ValueError(f"Unsupported device: {device_kind}")
 

brainstate/event/_csr_benchmark.py

@@ -0,0 +1,14 @@
+# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================