brainstate 0.1.0.post20250420__py2.py3-none-any.whl → 0.1.0.post20250423__py2.py3-none-any.whl

brainstate/nn/_dynamics/_dynamics_base_test.py

@@ -17,9 +17,10 @@
 
 from __future__ import annotations
 
-import numpy as np
 import unittest
 
+import numpy as np
+
 import brainstate as bst
 
 

brainstate/nn/_dynamics/_projection_base.py

@@ -104,35 +104,38 @@ class _AlignPost(Module):
 
 class AlignPostProj(Interaction):
     """
-    Full-chain synaptic projection with the align-post reduction and the automatic synapse merging.
-
-    The ``full-chain`` means that the model needs to provide all information needed for a projection,
-    including ``pre`` -> ``delay`` -> ``comm`` -> ``syn`` -> ``out`` -> ``post``.
-
-    The ``align-post`` means that the synaptic variables have the same dimension as the post-synaptic neuron group.
-
-    The ``merging`` means that the same delay model is shared by all synapses, and the synapse model with same
-    parameters (such like time constants) will also share the same synaptic variables.
-
-    All align-post projection models prefer to use the event-driven computation mode. This means that the
-    ``comm`` model should be the event-driven model.
-
-    Moreover, it's worth noting that ``FullProjAlignPost`` has a different updating order with all align-pre
-    projection models. The updating order of align-post projections is ``spikes`` -> ``comm`` -> ``syn`` -> ``out``.
-    While, the updating order of all align-pre projection models is usually ``spikes`` -> ``syn`` -> ``comm`` -> ``out``.
+    Align-post projection of the neural network.
+
+
+    Examples
+    --------
+
+    Here is an example of using the `AlignPostProj` to create a synaptic projection.
+    Note that this projection needs the manual input of pre-synaptic spikes.
+
+    >>> import brainstate
+    >>> import brainevent
+    >>> import brainunit as u
+    >>> n_exc = 3200
+    >>> n_inh = 800
+    >>> num = n_exc + n_inh
+    >>> pop = brainstate.nn.LIFRef(
+    ...        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=brainstate.init.Normal(-55., 2., unit=u.mV)
+    ... )
+    >>> pop.reset_state()
+    >>> E = brainstate.nn.AlignPostProj(
+    ...        comm=brainevent.nn.FixedProb(n_exc, num, prob=80 / num, weight=1.62 * u.mS),
+    ...        syn=brainstate.nn.Expon.desc(num, tau=5. * u.ms),
+    ...        out=brainstate.nn.CUBA.desc(scale=u.volt),
+    ...        post=pop
+    ... )
+    >>> exe_current = E(pop.spike.value)
 
 
-    #  brainstate.nn.AlignPostProj(
-    #       LIF().prefetch('V').delay.at('I'), bst.surrogate.ReluGrad(), comm, syn, out, post
-    #  )
 
-    Args:
-      pre: The pre-synaptic neuron group.
-      delay: The synaptic delay.
-      comm: The synaptic communication.
-      syn: The synaptic dynamics.
-      out: The synaptic output.
-      post: The post-synaptic neuron group.
     """
     __module__ = 'brainstate.nn'
 
@@ -314,10 +317,6 @@ class CurrentProj(Interaction):
     This is because the ``comm`` is computed after the synapse state, which is a floating-point number, rather
     than the spiking. To facilitate the event-driven computation, please use align post projections.
 
-    #    bint.CurrentInteraction(
-    #       LIF().align_pre(bst.nn.Expon.desc()).prefetch('g'), comm, out, post
-    #    )
-
     Args:
       prefetch: The synaptic dynamics.
       comm: The synaptic communication.

brainstate/nn/_dynamics/_state_delay.py

@@ -16,14 +16,14 @@
 from __future__ import annotations
 
 import math
+import numbers
+from functools import partial
+from typing import Optional, Dict, Callable, Union, Sequence
 
 import brainunit as u
 import jax
 import jax.numpy as jnp
-import numbers
 import numpy as np
-from functools import partial
-from typing import Optional, Dict, Callable, Union, Sequence
 
 from brainstate import environ
 from brainstate._state import ShortTermState, State

brainstate/nn/_dynamics/_synouts_test.py

@@ -15,10 +15,11 @@
 
 from __future__ import annotations
 
+import unittest
+
 import brainunit as u
 import jax.numpy as jnp
 import numpy as np
-import unittest
 
 import brainstate as bst
 

brainstate/nn/_elementwise/_dropout.py

@@ -16,11 +16,12 @@
 
 from __future__ import annotations
 
-import brainunit as u
-import jax.numpy as jnp
 from functools import partial
 from typing import Optional, Sequence
 
+import brainunit as u
+import jax.numpy as jnp
+
 from brainstate import random, environ, init
 from brainstate._state import ShortTermState
 from brainstate.nn._module import ElementWiseBlock

brainstate/nn/_elementwise/_dropout_test.py

@@ -14,9 +14,10 @@
 # ==============================================================================
 
 
-import numpy as np
 import unittest
 
+import numpy as np
+
 import brainstate as bst
 
 

brainstate/nn/_elementwise/_elementwise.py

@@ -17,10 +17,11 @@
 
 from __future__ import annotations
 
+from typing import Optional
+
 import brainunit as u
 import jax.numpy as jnp
 import jax.typing
-from typing import Optional
 
 from brainstate import random, functional as F
 from brainstate._state import ParamState

brainstate/nn/{_dyn_impl/_projection_alignpost.py → _event/__init__.py}

@@ -1,4 +1,4 @@
-# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
+# Copyright 2025 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.
@@ -13,13 +13,14 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
+# -*- coding: utf-8 -*-
 
-from brainstate.nn._dynamics._dynamics_base import Projection
+
+from ._fixedprob_mv import EventFixedProb, EventFixedNumConn
+from ._linear_mv import EventLinear
 
 __all__ = [
+    'EventLinear',
+    'EventFixedProb',
+    'EventFixedNumConn',
 ]
-
-
-class ExponentialSynapse(Projection):
-    pass

brainstate/nn/_event/_fixedprob_mv.py

@@ -0,0 +1,169 @@
+# 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.
+# ==============================================================================
+
+from __future__ import annotations
+
+from typing import Union, Callable, Optional
+
+import brainunit as u
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from brainstate import random, augment, environ, init
+from brainstate._compatible_import import brainevent
+from brainstate._state import ParamState
+from brainstate.compile import for_loop
+from brainstate.nn._module import Module
+from brainstate.typing import Size, ArrayLike
+
+__all__ = [
+    'EventFixedNumConn',
+    'EventFixedProb',
+]
+
+
+def init_indices_without_replace(
+    conn_num: int,
+    n_pre: int,
+    n_post: int,
+    seed: int | None,
+    method: str
+):
+    rng = random.default_rng(seed)
+
+    if method == 'vmap':
+        @augment.vmap
+        def rand_indices(key):
+            rng.set_key(key)
+            return rng.choice(n_post, size=(conn_num,), replace=False)
+
+        return rand_indices(rng.split_key(n_pre))
+
+    elif method == 'for_loop':
+        return for_loop(
+            lambda *args: rng.choice(n_post, size=(conn_num,), replace=False),
+            length=n_pre
+        )
+
+    else:
+        raise ValueError(f"Unknown method: {method}")
+
+
+class EventFixedNumConn(Module):
+    """
+    The FixedProb module implements a fixed probability connection with CSR sparse data structure.
+
+    Parameters
+    ----------
+    in_size : Size
+        Number of pre-synaptic neurons, i.e., input size.
+    out_size : Size
+        Number of post-synaptic neurons, i.e., output size.
+    conn_num : float, int
+        If it is a float, representing the probability of connection, i.e., connection probability.
+
+        If it is an integer, representing the number of connections.
+    conn_weight : float or callable or jax.Array or brainunit.Quantity
+        Maximum synaptic conductance, i.e., synaptic weight.
+    conn_target : str, optional
+        The target of the connection. Default is 'post', meaning that each pre-synaptic neuron connects to
+        a fixed number of post-synaptic neurons. The connection number is determined by the value of ``n_conn``.
+
+        If 'pre', each post-synaptic neuron connects to a fixed number of pre-synaptic neurons.
+    conn_init : str, optional
+        The initialization method of the connection weight. Default is 'vmap', meaning that the connection weight
+        is initialized by parallelized across multiple threads.
+
+        If 'for_loop', the connection weight is initialized by a for loop.
+    allow_multi_conn : bool, optional
+        Whether multiple connections are allowed from a single pre-synaptic neuron.
+        Default is True, meaning that a value of ``a`` can be selected multiple times.
+    seed: int, optional
+        Random seed. Default is None. If None, the default random seed will be used.
+    name : str, optional
+        Name of the module.
+    """
+
+    __module__ = 'brainstate.nn'
+
+    def __init__(
+        self,
+        in_size: Size,
+        out_size: Size,
+        conn_num: Union[int, float],
+        conn_weight: Union[Callable, ArrayLike],
+        conn_target: str = 'post',  # 'pre' or 'post'
+        allow_multi_conn: bool = True,
+        seed: Optional[int] = None,
+        name: Optional[str] = None,
+        conn_init: str = 'vmap',  # 'vmap' or 'for_loop'
+        param_type: type = ParamState,
+    ):
+        super().__init__(name=name)
+
+        # network parameters
+        self.in_size = in_size
+        self.out_size = out_size
+        self.conn_target = conn_target
+        assert conn_target in ('pre', 'post'), 'The target of the connection must be either "pre" or "post".'
+        if isinstance(conn_num, float):
+            assert 0. <= conn_num <= 1., 'Connection probability must be in [0, 1].'
+            conn_num = int(self.out_size[-1] * conn_num) if conn_target == 'post' else int(self.in_size[-1] * conn_num)
+        assert isinstance(conn_num, int), 'Connection number must be an integer.'
+        self.conn_num = conn_num
+        self.seed = seed
+        self.allow_multi_conn = allow_multi_conn
+
+        # connections
+        if self.conn_num >= 1:
+            if self.conn_target == 'post':
+                n_post = self.out_size[-1]
+                n_pre = self.in_size[-1]
+            else:
+                n_post = self.in_size[-1]
+                n_pre = self.out_size[-1]
+
+            # indices of post connected neurons
+            with jax.ensure_compile_time_eval():
+                if allow_multi_conn:
+                    rng = np.random if seed is None else np.random.RandomState(seed)
+                    indices = rng.randint(0, n_post, size=(n_pre, self.conn_num))
+                else:
+                    indices = init_indices_without_replace(self.conn_num, n_pre, n_post, seed, conn_init)
+                indices = u.math.asarray(indices, dtype=environ.ditype())
+            conn_weight = init.param(conn_weight, (n_pre, self.conn_num), allow_none=False)
+            conn_weight = u.math.asarray(conn_weight)
+            self.weight = param_type(conn_weight)
+            csr = (
+                brainevent.FixedPostNumConn((conn_weight, indices), shape=(n_pre, n_post))
+                if self.conn_target == 'post' else
+                brainevent.FixedPreNumConn((conn_weight, indices), shape=(n_pre, n_post))
+            )
+            self.conn = csr
+
+    def update(self, spk: jax.Array) -> Union[jax.Array, u.Quantity]:
+        if self.conn_num >= 1:
+            csr = self.conn.with_data(self.weight.value)
+            return brainevent.EventArray(spk) @ csr
+        else:
+            weight = self.weight.value
+            unit = u.get_unit(weight)
+            r = jnp.zeros(spk.shape[:-1] + (self.out_size[-1],), dtype=weight.dtype)
+            r = u.maybe_decimal(u.Quantity(r, unit=unit))
+            return u.math.asarray(r, dtype=environ.dftype())
+
+
+EventFixedProb = EventFixedNumConn

brainstate/nn/_event/_fixedprob_mv_test.py

@@ -0,0 +1,115 @@
+# 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.
+# ==============================================================================
+
+from __future__ import annotations
+
+import jax.numpy
+import jax.numpy as jnp
+import pytest
+
+import brainstate
+
+
+class TestFixedProbCSR:
+    @pytest.mark.parametrize('allow_multi_conn', [True, False, ])
+    def test1(self, allow_multi_conn):
+        x = brainstate.random.rand(20) < 0.1
+        # x = brainstate.random.rand(20)
+        m = brainstate.nn.EventFixedProb(20, 40, 0.1, 1.0, seed=123, allow_multi_conn=allow_multi_conn)
+        y = m(x)
+        print(y)
+
+        m2 = brainstate.nn.EventFixedProb(20, 40, 0.1, brainstate.init.KaimingUniform(), seed=123)
+        print(m2(x))
+
+    def test_grad_bool(self):
+        n_in = 20
+        n_out = 30
+        x = jax.numpy.asarray(brainstate.random.rand(n_in) < 0.3, dtype=float)
+        fn = brainstate.nn.EventFixedProb(n_in, n_out, 0.1, brainstate.init.KaimingUniform(), seed=123)
+
+        def f(x):
+            return fn(x).sum()
+
+        print(jax.grad(f)(x))
+
+    @pytest.mark.parametrize('homo_w', [True, False])
+    def test_vjp(self, homo_w):
+        n_in = 20
+        n_out = 30
+        x = jax.numpy.asarray(brainstate.random.rand(n_in) < 0.3, dtype=float)
+
+        if homo_w:
+            fn = brainstate.nn.EventFixedProb(n_in, n_out, 0.1, 1.5, seed=123)
+        else:
+            fn = brainstate.nn.EventFixedProb(n_in, n_out, 0.1, brainstate.init.KaimingUniform(), seed=123)
+        w = fn.weight.value
+
+        def f(x, w):
+            fn.weight.value = w
+            return fn(x).sum()
+
+        r = brainstate.augment.grad(f, argnums=(0, 1))(x, w)
+
+        # -------------------
+        # TRUE gradients
+
+        def true_fn(x, w, indices, n_post):
+            post = jnp.zeros((n_post,))
+            for i in range(n_in):
+                post = post.at[indices[i]].add(w * x[i] if homo_w else w[i] * x[i])
+            return post
+
+        def f2(x, w):
+            return true_fn(x, w, fn.conn.indices, n_out).sum()
+
+        r2 = jax.grad(f2, argnums=(0, 1))(x, w)
+        assert (jnp.allclose(r[0], r2[0]))
+        assert (jnp.allclose(r[1], r2[1]))
+
+    @pytest.mark.parametrize('homo_w', [True, False])
+    def test_jvp(self, homo_w):
+        n_in = 20
+        n_out = 30
+        x = jax.numpy.asarray(brainstate.random.rand(n_in) < 0.3, dtype=float)
+
+        fn = brainstate.nn.EventFixedProb(
+            n_in, n_out, 0.1, 1.5 if homo_w else brainstate.init.KaimingUniform(),
+            seed=123,
+        )
+        w = fn.weight.value
+
+        def f(x, w):
+            fn.weight.value = w
+            return fn(x)
+
+        o1, r1 = jax.jvp(f, (x, w), (jnp.ones_like(x), jnp.ones_like(w)))
+
+        # -------------------
+        # TRUE gradients
+
+        def true_fn(x, w, indices, n_post):
+            post = jnp.zeros((n_post,))
+            for i in range(n_in):
+                post = post.at[indices[i]].add(w * x[i] if homo_w else w[i] * x[i])
+            return post
+
+        def f2(x, w):
+            return true_fn(x, w, fn.conn.indices, n_out)
+
+        o2, r2 = jax.jvp(f2, (x, w), (jnp.ones_like(x), jnp.ones_like(w)))
+        assert (jnp.allclose(o1, o2))
+        # assert jnp.allclose(r1, r2), f'r1={r1}, r2={r2}'
+        assert (jnp.allclose(r1, r2, rtol=1e-4, atol=1e-4))

brainstate/nn/_event/_linear_mv.py

@@ -0,0 +1,85 @@
+# 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.
+# ==============================================================================
+
+from __future__ import annotations
+
+from typing import Union, Callable, Optional
+
+import brainunit as u
+import jax
+
+from brainstate import init
+from brainstate._compatible_import import brainevent
+from brainstate._state import ParamState
+from brainstate.nn._module import Module
+from brainstate.typing import Size, ArrayLike
+
+__all__ = [
+    'EventLinear',
+]
+
+
+class EventLinear(Module):
+    """
+
+    Parameters
+    ----------
+    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.
+    block_size : int, optional
+        Block size for parallel computation.
+    float_as_event : bool, optional
+        Whether to treat float as event.
+    name : str, optional
+        Name of the module.
+    """
+
+    __module__ = 'brainstate.nn'
+
+    def __init__(
+        self,
+        in_size: Size,
+        out_size: Size,
+        weight: Union[Callable, ArrayLike],
+        float_as_event: bool = True,
+        block_size: int = 64,
+        name: Optional[str] = None,
+        param_type: type = ParamState,
+    ):
+        super().__init__(name=name)
+
+        # network parameters
+        self.in_size = in_size
+        self.out_size = out_size
+        self.float_as_event = float_as_event
+        self.block_size = block_size
+
+        # maximum synaptic conductance
+        weight = init.param(weight, (self.in_size[-1], self.out_size[-1]), allow_none=False)
+        self.weight = param_type(weight)
+
+    def update(self, spk: jax.Array) -> Union[jax.Array, u.Quantity]:
+        weight = self.weight.value
+        if u.math.size(weight) == 1:
+            return u.math.ones(self.out_size) * (u.math.sum(spk) * weight)
+
+        if self.float_as_event:
+            return brainevent.EventArray(spk) @ weight
+        else:
+            return spk @ weight

brainstate/nn/_event/_linear_mv_test.py

@@ -0,0 +1,121 @@
+# 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.
+# ==============================================================================
+
+from __future__ import annotations
+
+import jax
+import jax.numpy as jnp
+import pytest
+
+import brainstate
+
+
+class TestEventLinear:
+    @pytest.mark.parametrize('bool_x', [True, False])
+    @pytest.mark.parametrize('homo_w', [True, False])
+    def test1(self, homo_w, bool_x):
+        x = brainstate.random.rand(20) < 0.1
+        if not bool_x:
+            x = jnp.asarray(x, dtype=float)
+        m = brainstate.nn.EventLinear(
+            20, 40,
+            1.5 if homo_w else brainstate.init.KaimingUniform(),
+            float_as_event=bool_x
+        )
+        y = m(x)
+        print(y)
+
+        assert (jnp.allclose(y, (x.sum() * m.weight.value) if homo_w else (x @ m.weight.value)))
+
+    def test_grad_bool(self):
+        n_in = 20
+        n_out = 30
+        x = brainstate.random.rand(n_in) < 0.3
+        fn = brainstate.nn.EventLinear(n_in, n_out, brainstate.init.KaimingUniform())
+
+        with pytest.raises(TypeError):
+            print(jax.grad(lambda x: fn(x).sum())(x))
+
+    @pytest.mark.parametrize('bool_x', [True, False])
+    @pytest.mark.parametrize('homo_w', [True, False])
+    def test_vjp(self, bool_x, homo_w):
+        n_in = 20
+        n_out = 30
+        if bool_x:
+            x = jax.numpy.asarray(brainstate.random.rand(n_in) < 0.3, dtype=float)
+        else:
+            x = brainstate.random.rand(n_in)
+
+        fn = brainstate.nn.EventLinear(
+            n_in,
+            n_out,
+            1.5 if homo_w else brainstate.init.KaimingUniform(),
+            float_as_event=bool_x
+        )
+        w = fn.weight.value
+
+        def f(x, w):
+            fn.weight.value = w
+            return fn(x).sum()
+
+        r1 = jax.grad(f, argnums=(0, 1))(x, w)
+
+        # -------------------
+        # TRUE gradients
+
+        def f2(x, w):
+            y = (x @ (jnp.ones([n_in, n_out]) * w)) if homo_w else (x @ w)
+            return y.sum()
+
+        r2 = jax.grad(f2, argnums=(0, 1))(x, w)
+        assert (jnp.allclose(r1[0], r2[0]))
+
+        if not jnp.allclose(r1[1], r2[1]):
+            print(r1[1] - r2[1])
+
+        assert (jnp.allclose(r1[1], r2[1]))
+
+    @pytest.mark.parametrize('bool_x', [True, False])
+    @pytest.mark.parametrize('homo_w', [True, False])
+    def test_jvp(self, bool_x, homo_w):
+        n_in = 20
+        n_out = 30
+        if bool_x:
+            x = jax.numpy.asarray(brainstate.random.rand(n_in) < 0.3, dtype=float)
+        else:
+            x = brainstate.random.rand(n_in)
+
+        fn = brainstate.nn.EventLinear(
+            n_in, n_out, 1.5 if homo_w else brainstate.init.KaimingUniform(),
+            float_as_event=bool_x
+        )
+        w = fn.weight.value
+
+        def f(x, w):
+            fn.weight.value = w
+            return fn(x)
+
+        o1, r1 = jax.jvp(f, (x, w), (jnp.ones_like(x), jnp.ones_like(w)))
+
+        # -------------------
+        # TRUE gradients
+
+        def f2(x, w):
+            y = (x @ (jnp.ones([n_in, n_out]) * w)) if homo_w else (x @ w)
+            return y
+
+        o2, r2 = jax.jvp(f, (x, w), (jnp.ones_like(x), jnp.ones_like(w)))
+        assert (jnp.allclose(o1, o2))
+        assert (jnp.allclose(r1, r2))