brainstate 0.1.3__py2.py3-none-any.whl → 0.1.4__py2.py3-none-any.whl

brainstate/nn/_dynamics.py

@@ -36,13 +36,14 @@ For handling the delays:
 from typing import Any, Dict, Callable, Hashable, Optional, Union, TypeVar, TYPE_CHECKING
 
 import brainunit as u
+import jax
 import numpy as np
 
 from brainstate import environ
 from brainstate._state import State
 from brainstate.graph import Node
 from brainstate.mixin import ParamDescriber, UpdateReturn
-from brainstate.typing import Size, ArrayLike
+from brainstate.typing import Size, ArrayLike, PyTree
 from ._delay import StateWithDelay, Delay
 from ._module import Module
 
@@ -811,18 +812,25 @@ class Dynamics(Module, UpdateReturn):
         >>> n1.align_pre(brainstate.nn.Expon.desc(n1.varshape))  # n2 will run after n1
         """
         if isinstance(dyn, Dynamics):
-            self._add_after_update(dyn.name, dyn)
+            self._add_after_update(id(dyn), dyn)
             return dyn
         elif isinstance(dyn, ParamDescriber):
             if not issubclass(dyn.cls, Dynamics):
                 raise TypeError(f'The input {dyn} should be an instance of {Dynamics}.')
             if not self._has_after_update(dyn.identifier):
-                self._add_after_update(dyn.identifier, dyn())
+                self._add_after_update(
+                    dyn.identifier,
+                    dyn() if ('in_size' in dyn.kwargs or len(dyn.args) > 0) else dyn(in_size=self.varshape)
+                )
             return self._get_after_update(dyn.identifier)
         else:
             raise TypeError(f'The input {dyn} should be an instance of {Dynamics} or a delayed initializer.')
 
-    def prefetch_delay(self, state: str, delay: Optional[ArrayLike] = None) -> 'PrefetchDelayAt':
+    def prefetch_delay(
+        self,
+        state: str,
+        delay: Optional[ArrayLike] = None
+    ) -> 'PrefetchDelayAt':
         """
         Create a reference to a delayed state or variable in the module.
 
@@ -840,7 +848,11 @@ class Dynamics(Module, UpdateReturn):
         """
         return self.prefetch(state).delay.at(delay)
 
-    def output_delay(self, delay: Optional[ArrayLike] = None) -> 'OutputDelayAt':
+    def output_delay(
+        self,
+        delay: Optional[ArrayLike] = None,
+        variable_like: PyTree = None
+    ) -> 'OutputDelayAt':
         """
         Create a reference to the delayed output of the module.
 
@@ -851,6 +863,7 @@ class Dynamics(Module, UpdateReturn):
         Args:
             delay (Optional[ArrayLike]): The amount of time to delay the output access,
                 typically in time units (e.g., milliseconds). Defaults to None.
+            variable_like:
 
         Returns:
             OutputDelayAt: An object that provides access to the module's output at the specified delay time.
@@ -1102,8 +1115,6 @@ class OutputDelayAt(Node):
     ----------
     module : Dynamics
         The dynamics module that contains the referenced state or variable.
-    item : str
-        The name of the state or variable to access with delay.
     time : ArrayLike
         The amount of time to delay access by, typically in time units (e.g., milliseconds).
 
@@ -1113,56 +1124,40 @@ class OutputDelayAt(Node):
     >>> import brainunit as u
     >>> neuron = brainstate.nn.LIF(10)
     >>> # Create a reference to voltage delayed by 5ms
-    >>> delayed_v = PrefetchDelayAt(neuron, 'V', 5.0 * u.ms)
+    >>> delayed_spike = OutputDelayAt(neuron, 5.0 * u.ms)
     >>> # Get the delayed value
-    >>> v_value = delayed_v()
+    >>> v_value = delayed_spike()
     """
 
     def __init__(
         self,
         module: Dynamics,
         time: Optional[ArrayLike] = None,
+        variable_like: Optional[PyTree] = None,
     ):
-        """
-        Initialize a PrefetchDelayAt object.
-
-        Parameters
-        ----------
-        module : AlignPre, Module
-            The dynamics module that contains the referenced state or variable.
-        time : ArrayLike
-            The amount of time to delay access by, typically in time units.
-        """
         super().__init__()
-        assert isinstance(module, UpdateReturn), 'The module should implement the `update_return` method.'
-        assert isinstance(module, Module), 'The module should be an instance of Module.'
+        assert isinstance(module, Dynamics), 'The module should be an instance of Dynamics.'
         self.module = module
+        dt = environ.get_dt()
+        if time is None:
+            time = u.math.zeros_like(dt)
         self.time = time
-        if time is not None:
-            self.step = u.math.asarray(time / environ.get_dt(), dtype=environ.ditype())
-
-            # register the delay
-            key = _get_output_delay_key()
-            if not module._has_after_update(key):
-                # TODO: unit processing
-                delay = Delay(module.update_return(), time)
-                module._add_after_update(key, receive_update_output(delay))
-            self.out_delay: Delay = module._get_after_update(key)
-            self.out_delay.register_delay(time)
+        self.step = u.math.asarray(time / dt, dtype=environ.ditype())
+
+        # register the delay
+        key = _get_output_delay_key()
+        if not module._has_after_update(key):
+            delay = Delay(
+                jax.ShapeDtypeStruct(module.out_size, dtype=environ.dftype()),
+                time,
+                take_aware_unit=True
+            )
+            module._add_after_update(key, receive_update_output(delay))
+        self.out_delay: Delay = module._get_after_update(key)
+        self.out_delay.register_delay(time)
 
     def __call__(self, *args, **kwargs):
-        """
-        Retrieve the value of the state at the specified delay time.
-
-        Returns
-        -------
-        Any
-            The value of the state or variable at the specified delay time.
-        """
-        if self.time is None:
-            return self.module.update_return()
-        else:
-            return self.out_delay.retrieve_at_step(self.step)
+        return self.out_delay.retrieve_at_step(self.step)
 
 
 def _get_prefetch_delay_key(item) -> str:

brainstate/nn/{_fixedprob_mv.py → _fixedprob.py}

@@ -16,19 +16,20 @@
 
 from typing import Union, Callable, Optional
 
+import brainevent
 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._state import ParamState, FakeState
 from brainstate.compile import for_loop
 from brainstate.typing import Size, ArrayLike
 from ._module import Module
 
 __all__ = [
+    'FixedNumConn',
     'EventFixedNumConn',
     'EventFixedProb',
 ]
@@ -44,12 +45,11 @@ def init_indices_without_replace(
     rng = random.default_rng(seed)
 
     if method == 'vmap':
-        @augment.vmap
-        def rand_indices(key):
-            rng.set_key(key)
+        @augment.vmap(axis_size=n_pre)
+        def rand_indices():
             return rng.choice(n_post, size=(conn_num,), replace=False)
 
-        return rand_indices(rng.split_key(n_pre))
+        return rand_indices()
 
     elif method == 'for_loop':
         return for_loop(
@@ -61,9 +61,9 @@ def init_indices_without_replace(
         raise ValueError(f"Unknown method: {method}")
 
 
-class EventFixedNumConn(Module):
+class FixedNumConn(Module):
     """
-    The FixedProb module implements a fixed probability connection with CSR sparse data structure.
+    The ``FixedNumConn`` module implements a fixed probability connection with CSR sparse data structure.
 
     Parameters
     ----------
@@ -77,7 +77,7 @@ class EventFixedNumConn(Module):
         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
+    efferent_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``.
 
@@ -104,7 +104,8 @@ class EventFixedNumConn(Module):
         out_size: Size,
         conn_num: Union[int, float],
         conn_weight: Union[Callable, ArrayLike],
-        conn_target: str = 'post',  # 'pre' or 'post'
+        efferent_target: str = 'post',  # 'pre' or 'post'
+        afferent_ratio: Union[int, float] = 1.,
         allow_multi_conn: bool = True,
         seed: Optional[int] = None,
         name: Optional[str] = None,
@@ -116,11 +117,14 @@ class EventFixedNumConn(Module):
         # 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".'
+        self.efferent_target = efferent_target
+        assert efferent_target in ('pre', 'post'), 'The target of the connection must be either "pre" or "post".'
+        assert 0. <= afferent_ratio <= 1., 'Afferent ratio must be in [0, 1].'
         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)
+            conn_num = (int(self.out_size[-1] * conn_num)
+                        if efferent_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
@@ -128,14 +132,13 @@ class EventFixedNumConn(Module):
 
         # connections
         if self.conn_num >= 1:
-            if self.conn_target == 'post':
+            if self.efferent_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)
@@ -143,15 +146,83 @@ class EventFixedNumConn(Module):
                 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
+
+            if afferent_ratio == 1.:
+                conn_weight = u.math.asarray(init.param(conn_weight, (n_pre, self.conn_num), allow_none=False))
+                self.weight = param_type(conn_weight)
+                csr = (
+                    brainevent.FixedPostNumConn((conn_weight, indices), shape=(n_pre, n_post))
+                    if self.efferent_target == 'post' else
+                    brainevent.FixedPreNumConn((conn_weight, indices), shape=(n_pre, n_post))
+                )
+                self.conn = csr
+
+            else:
+                self.pre_selected = np.random.random(n_pre) < afferent_ratio
+                indices = indices[self.pre_selected].flatten()
+                conn_weight = u.math.asarray(init.param(conn_weight, (indices.size,), allow_none=False))
+                self.weight = param_type(conn_weight)
+                indptr = (jnp.arange(1, n_pre + 1) * self.conn_num -
+                          jnp.cumsum(~self.pre_selected) * self.conn_num)
+                indptr = jnp.insert(indptr, 0, 0)  # insert 0 at the beginning
+                csr = (
+                    brainevent.CSR((conn_weight, indices, indptr), shape=(n_pre, n_post))
+                    if self.efferent_target == 'post' else
+                    brainevent.CSC((conn_weight, indices, indptr), shape=(n_pre, n_post))
+                )
+                self.conn = csr
+
+        else:
+            conn_weight = u.math.asarray(init.param(conn_weight, (), allow_none=False))
+            self.weight = FakeState(conn_weight)
+
+    def update(self, x: jax.Array) -> Union[jax.Array, u.Quantity]:
+        if self.conn_num >= 1:
+            csr = self.conn.with_data(self.weight.value)
+            return x @ csr
+        else:
+            weight = self.weight.value
+            r = u.math.zeros(x.shape[:-1] + (self.out_size[-1],), dtype=weight.dtype)
+            r = u.maybe_decimal(u.Quantity(r, unit=u.get_unit(weight)))
+            return u.math.asarray(r, dtype=environ.dftype())
+
+
+class EventFixedNumConn(FixedNumConn):
+    """
+    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 update(self, spk: jax.Array) -> Union[jax.Array, u.Quantity]:
         if self.conn_num >= 1:

brainstate/nn/_linear_mv.py

@@ -19,7 +19,7 @@ import brainunit as u
 import jax
 
 from brainstate import init
-from brainstate._compatible_import import brainevent
+import brainevent
 from brainstate._state import ParamState
 from brainstate.typing import Size, ArrayLike
 from ._module import Module

brainstate/nn/_module.py

@@ -34,7 +34,7 @@ import numpy as np
 from brainstate._state import State
 from brainstate.graph import Node, states, nodes, flatten
 from brainstate.mixin import ParamDescriber, ParamDesc
-from brainstate.typing import PathParts
+from brainstate.typing import PathParts, Size
 from brainstate.util import FlattedDict, NestedDict, BrainStateError
 
 # maximum integer
@@ -62,8 +62,8 @@ class Module(Node, ParamDesc):
 
     __module__ = 'brainstate.nn'
 
-    _in_size: Optional[Tuple[int, ...]]
-    _out_size: Optional[Tuple[int, ...]]
+    _in_size: Optional[Size]
+    _out_size: Optional[Size]
     _name: Optional[str]
 
     if not TYPE_CHECKING:
@@ -87,7 +87,7 @@ class Module(Node, ParamDesc):
         raise AttributeError('The name of the model is read-only.')
 
     @property
-    def in_size(self) -> Tuple[int, ...]:
+    def in_size(self) -> Size:
         return self._in_size
 
     @in_size.setter
@@ -98,7 +98,7 @@ class Module(Node, ParamDesc):
         self._in_size = tuple(in_size)
 
     @property
-    def out_size(self) -> Tuple[int, ...]:
+    def out_size(self) -> Size:
         return self._out_size
 
     @out_size.setter