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

brainstate/compile/_make_jaxpr_test.py

@@ -88,7 +88,7 @@ class TestMakeJaxpr(unittest.TestCase):
         self.assertTrue(jnp.allclose(jaxpr_as_fun(jaxpr)(jnp.zeros(1), st1.value)[0],
                                      f3(jnp.zeros(1))))
 
-    def test_compar_jax_make_jaxpr2(self):
+    def test_compare_jax_make_jaxpr2(self):
         st1 = brainstate.State(jnp.ones(10))
 
         def fa(x):
@@ -108,7 +108,7 @@ class TestMakeJaxpr(unittest.TestCase):
         print(jaxpr)
         print(jaxpr_as_fun(jaxpr)(jnp.zeros(1)))
 
-    def test_compar_jax_make_jaxpr3(self):
+    def test_compare_jax_make_jaxpr3(self):
         def fa(x):
             return 1.
 
@@ -121,6 +121,17 @@ class TestMakeJaxpr(unittest.TestCase):
         print(jaxpr)
         # print(jaxpr_as_fun(jaxpr)(jnp.zeros(1)))
 
+    def test_static_argnames(self):
+        def func4(a, b):  # Arg is a pair
+            temp = a + jnp.sin(b) * 3.
+            c = brainstate.random.rand_like(a)
+            return jnp.sum(temp + c)
+
+        jaxpr, states = brainstate.compile.make_jaxpr(func4, static_argnames='b')(jnp.zeros(8), 1.)
+        print()
+        print(jaxpr)
+        print(states)
+
 
 def test_return_states():
     import jax.numpy

brainstate/graph/_graph_node.py

@@ -25,7 +25,7 @@ import numpy as np
 
 from brainstate._state import State, TreefyState
 from brainstate.typing import Key
-from brainstate.util._pretty_pytree import PrettyObject
+from brainstate.util.pretty_pytree import PrettyObject
 from ._graph_operation import register_graph_node_type
 
 __all__ = [

brainstate/graph/_graph_operation.py

@@ -30,10 +30,10 @@ from typing_extensions import TypeGuard, Unpack
 from brainstate._state import State, TreefyState
 from brainstate._utils import set_module_as
 from brainstate.typing import PathParts, Filter, Predicate, Key
-from brainstate.util._caller import ApplyCaller, CallableProxy, DelayedAccessor
-from brainstate.util._pretty_pytree import NestedDict, FlattedDict, PrettyDict
-from brainstate.util._pretty_repr import PrettyRepr, PrettyType, PrettyAttr, PrettyMapping, MappingReprMixin
-from brainstate.util._struct import FrozenDict
+from brainstate.util.caller import ApplyCaller, CallableProxy, DelayedAccessor
+from brainstate.util.pretty_pytree import NestedDict, FlattedDict, PrettyDict
+from brainstate.util.pretty_repr import PrettyRepr, PrettyType, PrettyAttr, PrettyMapping, MappingReprMixin
+from brainstate.util.struct import FrozenDict
 from brainstate.util.filter import to_predicate
 
 _max_int = np.iinfo(np.int32).max

brainstate/mixin.py

@@ -41,6 +41,14 @@ __all__ = [
 ]
 
 
+def hashable(x):
+    try:
+        hash(x)
+        return True
+    except TypeError:
+        return False
+
+
 class Mixin(object):
     """Base Mixin object.
 
@@ -67,6 +75,14 @@ class ParamDesc(Mixin):
 
 
 class HashableDict(dict):
+    def __init__(self, the_dict: dict):
+        out = dict()
+        for k, v in the_dict.items():
+            if not hashable(v):
+                v = str(v)  # convert to string if not hashable
+            out[k] = v
+        super().__init__(out)
+
     def __hash__(self):
         return hash(tuple(sorted(self.items())))
 
@@ -132,7 +148,6 @@ class AlignPost(Mixin):
         raise NotImplementedError
 
 
-
 class BindCondData(Mixin):
     """Bind temporary conductance data.
 
@@ -147,12 +162,26 @@ class BindCondData(Mixin):
         self._conductance = None
 
 
+def not_implemented(func):
+
+    def wrapper(*args, **kwargs):
+        raise NotImplementedError(f'{func.__name__} is not implemented.')
+
+    wrapper.not_implemented = True
+    return wrapper
+
+
+
 class UpdateReturn(Mixin):
+    @not_implemented
     def update_return(self) -> PyTree:
         """
         The update function return of the model.
 
-        It should be a pytree, with each element as a ``jax.ShapeDtypeStruct`` or ``jax.core.ShapedArray``.
+        This function requires no parameters and must return a PyTree.
+
+        It is usually used for delay initialization, for example, ``Dynamics.output_delay`` relies on this function to
+        initialize the output delay.
 
         """
         raise NotImplementedError(f'Must implement the "{self.update_return.__name__}()" function.')

brainstate/nn/__init__.py

@@ -33,12 +33,14 @@ from ._embedding import *
 from ._embedding import __all__ as embed_all
 from ._exp_euler import *
 from ._exp_euler import __all__ as exp_euler_all
-from ._fixedprob_mv import EventFixedProb, EventFixedNumConn
+from ._fixedprob import *
+from._fixedprob import __all__ as fixedprob_all
 from ._inputs import *
 from ._inputs import __all__ as inputs_all
 from ._linear import *
 from ._linear import __all__ as linear_all
-from ._linear_mv import EventLinear
+from ._linear_mv import *
+from ._linear_mv import __all__ as linear_mv_all
 from ._ltp import *
 from ._ltp import __all__ as ltp_all
 from ._module import *
@@ -69,9 +71,6 @@ from ._utils import __all__ as utils_all
 __all__ = (
     [
         'metrics',
-        'EventLinear',
-        'EventFixedProb',
-        'EventFixedNumConn',
     ]
     + collective_ops_all
     + common_all
@@ -87,6 +86,8 @@ __all__ = (
     + linear_all
     + normalizations_all
     + poolings_all
+    + fixedprob_all
+    + linear_mv_all
     + embed_all
     + dropout_all
     + elementwise_all
@@ -115,6 +116,8 @@ del (
     normalizations_all,
     poolings_all,
     embed_all,
+    fixedprob_all,
+    linear_mv_all,
     dropout_all,
     elementwise_all,
     dyn_neuron_all,

brainstate/nn/_common.py

@@ -118,14 +118,14 @@ class Vmap(Module):
     This class wraps a module and applies vectorized mapping to its execution,
     allowing for efficient parallel processing across specified axes.
 
-    Attributes:
+    Args:
         module (Module): The module to be vmapped.
-        in_axes (int | None | Sequence[Any]): Specifies how to map over inputs.
-        out_axes (Any): Specifies how to map over outputs.
-        vmap_states (Filter | Dict[Filter, int]): Specifies which states to vmap and on which axes.
-        vmap_out_states (Filter | Dict[Filter, int]): Specifies which output states to vmap and on which axes.
-        axis_name (AxisName | None): Name of the axis being mapped over.
-        axis_size (int | None): Size of the axis being mapped over.
+        in_axes (int | None | Sequence[Any], optional): Specifies how to map over inputs. Defaults to 0.
+        out_axes (Any, optional): Specifies how to map over outputs. Defaults to 0.
+        vmap_states (Filter | Dict[Filter, int], optional): Specifies which states to vmap and on which axes. Defaults to None.
+        vmap_out_states (Filter | Dict[Filter, int], optional): Specifies which output states to vmap and on which axes. Defaults to None.
+        axis_name (AxisName | None, optional): Name of the axis being mapped over. Defaults to None.
+        axis_size (int | None, optional): Size of the axis being mapped over. Defaults to None.
     """
 
     def __init__(
@@ -138,18 +138,6 @@ class Vmap(Module):
         axis_name: AxisName | None = None,
         axis_size: int | None = None,
     ):
-        """
-        Initialize the Vmap instance.
-
-        Args:
-            module (Module): The module to be vmapped.
-            in_axes (int | None | Sequence[Any], optional): Specifies how to map over inputs. Defaults to 0.
-            out_axes (Any, optional): Specifies how to map over outputs. Defaults to 0.
-            vmap_states (Filter | Dict[Filter, int], optional): Specifies which states to vmap and on which axes. Defaults to None.
-            vmap_out_states (Filter | Dict[Filter, int], optional): Specifies which output states to vmap and on which axes. Defaults to None.
-            axis_name (AxisName | None, optional): Name of the axis being mapped over. Defaults to None.
-            axis_size (int | None, optional): Size of the axis being mapped over. Defaults to None.
-        """
         super().__init__()
 
         # parameters

brainstate/nn/_delay.py

@@ -330,7 +330,14 @@ class Delay(Module):
         indices = (delay_idx,) + indices
 
         # the delay data
-        return jax.tree.map(lambda a: a[indices], self.history.value)
+        if self._unit is None:
+            return jax.tree.map(lambda a: a[indices], self.history.value)
+        else:
+            return jax.tree.map(
+                lambda hist, unit: u.maybe_decimal(hist[indices] * unit),
+                self.history.value,
+                self._unit
+            )
 
     def retrieve_at_time(self, delay_time, *indices) -> PyTree:
         """
@@ -393,6 +400,9 @@ class Delay(Module):
         """
         assert self.history is not None, 'The delay history is not initialized.'
 
+        if self.take_aware_unit and self._unit is None:
+            self._unit = jax.tree.map(lambda x: u.get_unit(x), current, is_leaf=u.math.is_quantity)
+
         # update the delay data at the rotation index
         if self.delay_method == _DELAY_ROTATE:
             i = environ.get(environ.I)
@@ -419,6 +429,8 @@ class Delay(Module):
             raise ValueError(f'Unknown updating method "{self.delay_method}"')
 
 
+
+
 class StateWithDelay(Delay):
     """
     A ``State`` type that defines the state in a differential equation.

brainstate/nn/_dropout.py

@@ -409,7 +409,8 @@ class DropoutFixed(ElementWiseBlock):
         self.out_size = in_size
 
     def init_state(self, batch_size=None, **kwargs):
-        self.mask = ShortTermState(init.param(partial(random.bernoulli, self.prob), self.in_size, batch_size))
+        if self.prob < 1.:
+            self.mask = ShortTermState(init.param(partial(random.bernoulli, self.prob), self.in_size, batch_size))
 
     def update(self, x):
         dtype = u.math.get_dtype(x)
@@ -418,8 +419,8 @@ class DropoutFixed(ElementWiseBlock):
             if self.mask.value.shape != x.shape:
                 raise ValueError(f"Input shape {x.shape} does not match the mask shape {self.mask.value.shape}. "
                                  f"Please call `init_state()` method first.")
-            return jnp.where(self.mask.value,
-                             jnp.asarray(x / self.prob, dtype=dtype),
-                             jnp.asarray(0., dtype=dtype))
+            return u.math.where(self.mask.value,
+                                u.math.asarray(x / self.prob, dtype=dtype),
+                                u.math.asarray(0., dtype=dtype) * u.get_unit(x))
         else:
             return x

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/_exp_euler.py

@@ -69,27 +69,24 @@ def exp_euler_step(
             f'The input data type should be float64, float32, float16, or bfloat16 '
             f'when using Exponential Euler method. But we got {args[0].dtype}.'
         )
+
+    # drift
     dt = environ.get('dt')
     linear, derivative = vector_grad(fn, argnums=0, return_value=True)(*args, **kwargs)
     linear = u.Quantity(u.get_mantissa(linear), u.get_unit(derivative) / u.get_unit(linear))
     phi = u.math.exprel(dt * linear)
     x_next = args[0] + dt * phi * derivative
 
+    # diffusion
     if diffusion is not None:
-        # unit checking
-        diffusion = diffusion(*args, **kwargs)
-        time_unit = u.get_unit(dt)
-        drift_unit = u.get_unit(derivative)
-        diffusion_unit = u.get_unit(diffusion)
-        # if drift_unit.is_unitless:
-        #   assert diffusion_unit.is_unitless, 'The diffusion term should be unitless when the drift term is unitless.'
-        # else:
-        #   u.fail_for_dimension_mismatch(
-        #     drift_unit, diffusion_unit * time_unit ** 0.5,
-        #     "Drift unit is {drift}, diffusion unit is {diffusion}, ",
-        #     drift=drift_unit, diffusion=diffusion_unit * time_unit ** 0.5
-        #   )
-
-        # diffusion
-        x_next += diffusion * u.math.sqrt(dt) * random.randn_like(args[0])
+        diffusion_part = diffusion(*args, **kwargs) * u.math.sqrt(dt) * random.randn_like(args[0])
+        if u.get_dim(x_next) != u.get_dim(diffusion_part):
+            drift_unit = u.get_unit(x_next)
+            time_unit = u.get_unit(dt)
+            raise ValueError(
+                f"Drift unit is {drift_unit}, "
+                f"expected diffusion unit is {drift_unit / time_unit ** 0.5}, "
+                f"but we got {u.get_unit(diffusion_part)}."
+            )
+        x_next += diffusion_part
     return x_next

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/_inputs.py

@@ -547,7 +547,7 @@ def poisson_input(
                 num_input,
                 p,
                 tar[indices].shape,
-                # check_valid=False,
+                check_valid=False,
                 dtype=tar.dtype
             ),
             tar_val,