brainstate 0.2.0__py2.py3-none-any.whl → 0.2.1__py2.py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (112) hide show
  1. brainstate/__init__.py +169 -169
  2. brainstate/_compatible_import.py +340 -340
  3. brainstate/_compatible_import_test.py +681 -681
  4. brainstate/_deprecation.py +210 -210
  5. brainstate/_deprecation_test.py +2319 -2319
  6. brainstate/_error.py +45 -45
  7. brainstate/_state.py +1652 -1652
  8. brainstate/_state_test.py +52 -52
  9. brainstate/_utils.py +47 -47
  10. brainstate/environ.py +1495 -1495
  11. brainstate/environ_test.py +1223 -1223
  12. brainstate/graph/__init__.py +22 -22
  13. brainstate/graph/_node.py +240 -240
  14. brainstate/graph/_node_test.py +589 -589
  15. brainstate/graph/_operation.py +1624 -1624
  16. brainstate/graph/_operation_test.py +1147 -1147
  17. brainstate/mixin.py +1433 -1433
  18. brainstate/mixin_test.py +1017 -1017
  19. brainstate/nn/__init__.py +137 -137
  20. brainstate/nn/_activations.py +1100 -1100
  21. brainstate/nn/_activations_test.py +354 -354
  22. brainstate/nn/_collective_ops.py +633 -633
  23. brainstate/nn/_collective_ops_test.py +774 -774
  24. brainstate/nn/_common.py +226 -226
  25. brainstate/nn/_common_test.py +154 -154
  26. brainstate/nn/_conv.py +2010 -2010
  27. brainstate/nn/_conv_test.py +849 -849
  28. brainstate/nn/_delay.py +575 -575
  29. brainstate/nn/_delay_test.py +243 -243
  30. brainstate/nn/_dropout.py +618 -618
  31. brainstate/nn/_dropout_test.py +477 -477
  32. brainstate/nn/_dynamics.py +1267 -1267
  33. brainstate/nn/_dynamics_test.py +67 -67
  34. brainstate/nn/_elementwise.py +1298 -1298
  35. brainstate/nn/_elementwise_test.py +829 -829
  36. brainstate/nn/_embedding.py +408 -408
  37. brainstate/nn/_embedding_test.py +156 -156
  38. brainstate/nn/_event_fixedprob.py +233 -233
  39. brainstate/nn/_event_fixedprob_test.py +115 -115
  40. brainstate/nn/_event_linear.py +83 -83
  41. brainstate/nn/_event_linear_test.py +121 -121
  42. brainstate/nn/_exp_euler.py +254 -254
  43. brainstate/nn/_exp_euler_test.py +377 -377
  44. brainstate/nn/_linear.py +744 -744
  45. brainstate/nn/_linear_test.py +475 -475
  46. brainstate/nn/_metrics.py +1070 -1070
  47. brainstate/nn/_metrics_test.py +611 -611
  48. brainstate/nn/_module.py +384 -384
  49. brainstate/nn/_module_test.py +40 -40
  50. brainstate/nn/_normalizations.py +1334 -1334
  51. brainstate/nn/_normalizations_test.py +699 -699
  52. brainstate/nn/_paddings.py +1020 -1020
  53. brainstate/nn/_paddings_test.py +722 -722
  54. brainstate/nn/_poolings.py +2239 -2239
  55. brainstate/nn/_poolings_test.py +952 -952
  56. brainstate/nn/_rnns.py +946 -946
  57. brainstate/nn/_rnns_test.py +592 -592
  58. brainstate/nn/_utils.py +216 -216
  59. brainstate/nn/_utils_test.py +401 -401
  60. brainstate/nn/init.py +809 -809
  61. brainstate/nn/init_test.py +180 -180
  62. brainstate/random/__init__.py +270 -270
  63. brainstate/random/_rand_funs.py +3938 -3938
  64. brainstate/random/_rand_funs_test.py +640 -640
  65. brainstate/random/_rand_seed.py +675 -675
  66. brainstate/random/_rand_seed_test.py +48 -48
  67. brainstate/random/_rand_state.py +1617 -1617
  68. brainstate/random/_rand_state_test.py +551 -551
  69. brainstate/transform/__init__.py +59 -59
  70. brainstate/transform/_ad_checkpoint.py +176 -176
  71. brainstate/transform/_ad_checkpoint_test.py +49 -49
  72. brainstate/transform/_autograd.py +1025 -1025
  73. brainstate/transform/_autograd_test.py +1289 -1289
  74. brainstate/transform/_conditions.py +316 -316
  75. brainstate/transform/_conditions_test.py +220 -220
  76. brainstate/transform/_error_if.py +94 -94
  77. brainstate/transform/_error_if_test.py +52 -52
  78. brainstate/transform/_eval_shape.py +145 -145
  79. brainstate/transform/_eval_shape_test.py +38 -38
  80. brainstate/transform/_jit.py +399 -399
  81. brainstate/transform/_jit_test.py +143 -143
  82. brainstate/transform/_loop_collect_return.py +675 -675
  83. brainstate/transform/_loop_collect_return_test.py +58 -58
  84. brainstate/transform/_loop_no_collection.py +283 -283
  85. brainstate/transform/_loop_no_collection_test.py +50 -50
  86. brainstate/transform/_make_jaxpr.py +2016 -2016
  87. brainstate/transform/_make_jaxpr_test.py +1510 -1510
  88. brainstate/transform/_mapping.py +529 -529
  89. brainstate/transform/_mapping_test.py +194 -194
  90. brainstate/transform/_progress_bar.py +255 -255
  91. brainstate/transform/_random.py +171 -171
  92. brainstate/transform/_unvmap.py +256 -256
  93. brainstate/transform/_util.py +286 -286
  94. brainstate/typing.py +837 -837
  95. brainstate/typing_test.py +780 -780
  96. brainstate/util/__init__.py +27 -27
  97. brainstate/util/_others.py +1024 -1024
  98. brainstate/util/_others_test.py +962 -962
  99. brainstate/util/_pretty_pytree.py +1301 -1301
  100. brainstate/util/_pretty_pytree_test.py +675 -675
  101. brainstate/util/_pretty_repr.py +462 -462
  102. brainstate/util/_pretty_repr_test.py +696 -696
  103. brainstate/util/filter.py +945 -945
  104. brainstate/util/filter_test.py +911 -911
  105. brainstate/util/struct.py +910 -910
  106. brainstate/util/struct_test.py +602 -602
  107. {brainstate-0.2.0.dist-info → brainstate-0.2.1.dist-info}/METADATA +108 -108
  108. brainstate-0.2.1.dist-info/RECORD +111 -0
  109. {brainstate-0.2.0.dist-info → brainstate-0.2.1.dist-info}/licenses/LICENSE +202 -202
  110. brainstate-0.2.0.dist-info/RECORD +0 -111
  111. {brainstate-0.2.0.dist-info → brainstate-0.2.1.dist-info}/WHEEL +0 -0
  112. {brainstate-0.2.0.dist-info → brainstate-0.2.1.dist-info}/top_level.txt +0 -0
@@ -1,233 +1,233 @@
1
- # Copyright 2024 BrainX Ecosystem Limited. All Rights Reserved.
2
- #
3
- # Licensed under the Apache License, Version 2.0 (the "License");
4
- # you may not use this file except in compliance with the License.
5
- # You may obtain a copy of the License at
6
- #
7
- # http://www.apache.org/licenses/LICENSE-2.0
8
- #
9
- # Unless required by applicable law or agreed to in writing, software
10
- # distributed under the License is distributed on an "AS IS" BASIS,
11
- # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
- # See the License for the specific language governing permissions and
13
- # limitations under the License.
14
- # ==============================================================================
15
-
16
-
17
- from typing import Union, Callable, Optional
18
-
19
- import brainevent
20
- import brainunit as u
21
- import jax
22
- import jax.numpy as jnp
23
- import numpy as np
24
-
25
- from brainstate import random, transform, environ
26
- from brainstate._state import ParamState, FakeState
27
- from brainstate.transform import for_loop
28
- from brainstate.typing import Size, ArrayLike
29
- from . import init as init
30
- from ._module import Module
31
-
32
- __all__ = [
33
- 'FixedNumConn',
34
- 'EventFixedNumConn',
35
- 'EventFixedProb',
36
- ]
37
-
38
-
39
- def init_indices_without_replace(
40
- conn_num: int,
41
- n_pre: int,
42
- n_post: int,
43
- seed: int | None,
44
- method: str
45
- ):
46
- rng = random.default_rng(seed)
47
-
48
- if method == 'vmap':
49
- @transform.vmap(axis_size=n_pre)
50
- def rand_indices():
51
- return rng.choice(n_post, size=(conn_num,), replace=False)
52
-
53
- return rand_indices()
54
-
55
- elif method == 'for_loop':
56
- return for_loop(
57
- lambda *args: rng.choice(n_post, size=(conn_num,), replace=False),
58
- length=n_pre
59
- )
60
-
61
- else:
62
- raise ValueError(f"Unknown method: {method}")
63
-
64
-
65
- class FixedNumConn(Module):
66
- """
67
- The ``FixedNumConn`` module implements a fixed probability connection with CSR sparse data structure.
68
-
69
- Parameters
70
- ----------
71
- in_size : Size
72
- Number of pre-synaptic neurons, i.e., input size.
73
- out_size : Size
74
- Number of post-synaptic neurons, i.e., output size.
75
- conn_num : float, int
76
- If it is a float, representing the probability of connection, i.e., connection probability.
77
-
78
- If it is an integer, representing the number of connections.
79
- conn_weight : float or callable or jax.Array or brainunit.Quantity
80
- Maximum synaptic conductance, i.e., synaptic weight.
81
- efferent_target : str, optional
82
- The target of the connection. Default is 'post', meaning that each pre-synaptic neuron connects to
83
- a fixed number of post-synaptic neurons. The connection number is determined by the value of ``n_conn``.
84
-
85
- If 'pre', each post-synaptic neuron connects to a fixed number of pre-synaptic neurons.
86
- conn_init : str, optional
87
- The initialization method of the connection weight. Default is 'vmap', meaning that the connection weight
88
- is initialized by parallelized across multiple threads.
89
-
90
- If 'for_loop', the connection weight is initialized by a for loop.
91
- allow_multi_conn : bool, optional
92
- Whether multiple connections are allowed from a single pre-synaptic neuron.
93
- Default is True, meaning that a value of ``a`` can be selected multiple times.
94
- seed: int, optional
95
- Random seed. Default is None. If None, the default random seed will be used.
96
- name : str, optional
97
- Name of the module.
98
- """
99
-
100
- __module__ = 'brainstate.nn'
101
-
102
- def __init__(
103
- self,
104
- in_size: Size,
105
- out_size: Size,
106
- conn_num: Union[int, float],
107
- conn_weight: Union[Callable, ArrayLike],
108
- efferent_target: str = 'post', # 'pre' or 'post'
109
- afferent_ratio: Union[int, float] = 1.,
110
- allow_multi_conn: bool = True,
111
- seed: Optional[int] = None,
112
- name: Optional[str] = None,
113
- conn_init: str = 'vmap', # 'vmap' or 'for_loop'
114
- param_type: type = ParamState,
115
- ):
116
- super().__init__(name=name)
117
-
118
- # network parameters
119
- self.in_size = in_size
120
- self.out_size = out_size
121
- self.efferent_target = efferent_target
122
- assert efferent_target in ('pre', 'post'), 'The target of the connection must be either "pre" or "post".'
123
- assert 0. <= afferent_ratio <= 1., 'Afferent ratio must be in [0, 1].'
124
- if isinstance(conn_num, float):
125
- assert 0. <= conn_num <= 1., 'Connection probability must be in [0, 1].'
126
- conn_num = (int(self.out_size[-1] * conn_num)
127
- if efferent_target == 'post' else
128
- int(self.in_size[-1] * conn_num))
129
- assert isinstance(conn_num, int), 'Connection number must be an integer.'
130
- self.conn_num = conn_num
131
- self.seed = seed
132
- self.allow_multi_conn = allow_multi_conn
133
-
134
- # connections
135
- if self.conn_num >= 1:
136
- if self.efferent_target == 'post':
137
- n_post = self.out_size[-1]
138
- n_pre = self.in_size[-1]
139
- else:
140
- n_post = self.in_size[-1]
141
- n_pre = self.out_size[-1]
142
-
143
- with jax.ensure_compile_time_eval():
144
- if allow_multi_conn:
145
- rng = np.random if seed is None else np.random.RandomState(seed)
146
- indices = rng.randint(0, n_post, size=(n_pre, self.conn_num))
147
- else:
148
- indices = init_indices_without_replace(self.conn_num, n_pre, n_post, seed, conn_init)
149
- indices = u.math.asarray(indices, dtype=environ.ditype())
150
-
151
- if afferent_ratio == 1.:
152
- conn_weight = u.math.asarray(init.param(conn_weight, (n_pre, self.conn_num), allow_none=False))
153
- self.weight = param_type(conn_weight)
154
- csr = (
155
- brainevent.FixedPostNumConn((conn_weight, indices), shape=(n_pre, n_post))
156
- if self.efferent_target == 'post' else
157
- brainevent.FixedPreNumConn((conn_weight, indices), shape=(n_pre, n_post))
158
- )
159
- self.conn = csr
160
-
161
- else:
162
- self.pre_selected = np.random.random(n_pre) < afferent_ratio
163
- indices = indices[self.pre_selected].flatten()
164
- conn_weight = u.math.asarray(init.param(conn_weight, (indices.size,), allow_none=False))
165
- self.weight = param_type(conn_weight)
166
- indptr = (jnp.arange(1, n_pre + 1) * self.conn_num -
167
- jnp.cumsum(~self.pre_selected) * self.conn_num)
168
- indptr = jnp.insert(indptr, 0, 0) # insert 0 at the beginning
169
- csr = (
170
- brainevent.CSR((conn_weight, indices, indptr), shape=(n_pre, n_post))
171
- if self.efferent_target == 'post' else
172
- brainevent.CSC((conn_weight, indices, indptr), shape=(n_pre, n_post))
173
- )
174
- self.conn = csr
175
-
176
- else:
177
- conn_weight = u.math.asarray(init.param(conn_weight, (), allow_none=False))
178
- self.weight = FakeState(conn_weight)
179
-
180
- def update(self, x) -> Union[jax.Array, u.Quantity]:
181
- if self.conn_num >= 1:
182
- csr = self.conn.with_data(self.weight.value)
183
- return x @ csr
184
- else:
185
- weight = self.weight.value
186
- r = u.math.zeros(x.shape[:-1] + (self.out_size[-1],), dtype=weight.dtype)
187
- return u.maybe_decimal(u.Quantity(r, unit=u.get_unit(weight), dtype=environ.dftype()))
188
-
189
-
190
- class EventFixedNumConn(FixedNumConn):
191
- """
192
- The FixedProb module implements a fixed probability connection with CSR sparse data structure.
193
-
194
- Parameters
195
- ----------
196
- in_size : Size
197
- Number of pre-synaptic neurons, i.e., input size.
198
- out_size : Size
199
- Number of post-synaptic neurons, i.e., output size.
200
- conn_num : float, int
201
- If it is a float, representing the probability of connection, i.e., connection probability.
202
-
203
- If it is an integer, representing the number of connections.
204
- conn_weight : float or callable or jax.Array or brainunit.Quantity
205
- Maximum synaptic conductance, i.e., synaptic weight.
206
- conn_target : str, optional
207
- The target of the connection. Default is 'post', meaning that each pre-synaptic neuron connects to
208
- a fixed number of post-synaptic neurons. The connection number is determined by the value of ``n_conn``.
209
-
210
- If 'pre', each post-synaptic neuron connects to a fixed number of pre-synaptic neurons.
211
- conn_init : str, optional
212
- The initialization method of the connection weight. Default is 'vmap', meaning that the connection weight
213
- is initialized by parallelized across multiple threads.
214
-
215
- If 'for_loop', the connection weight is initialized by a for loop.
216
- allow_multi_conn : bool, optional
217
- Whether multiple connections are allowed from a single pre-synaptic neuron.
218
- Default is True, meaning that a value of ``a`` can be selected multiple times.
219
- seed: int, optional
220
- Random seed. Default is None. If None, the default random seed will be used.
221
- name : str, optional
222
- Name of the module.
223
- """
224
-
225
- __module__ = 'brainstate.nn'
226
-
227
- def update(self, spk: jax.Array) -> Union[jax.Array, u.Quantity]:
228
- return super().update(
229
- brainevent.EventArray(spk)
230
- )
231
-
232
-
233
- EventFixedProb = EventFixedNumConn
1
+ # Copyright 2024 BrainX Ecosystem Limited. All Rights Reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+ # ==============================================================================
15
+
16
+
17
+ from typing import Union, Callable, Optional
18
+
19
+ import brainevent
20
+ import brainunit as u
21
+ import jax
22
+ import jax.numpy as jnp
23
+ import numpy as np
24
+
25
+ from brainstate import random, transform, environ
26
+ from brainstate._state import ParamState, FakeState
27
+ from brainstate.transform import for_loop
28
+ from brainstate.typing import Size, ArrayLike
29
+ from . import init as init
30
+ from ._module import Module
31
+
32
+ __all__ = [
33
+ 'FixedNumConn',
34
+ 'EventFixedNumConn',
35
+ 'EventFixedProb',
36
+ ]
37
+
38
+
39
+ def init_indices_without_replace(
40
+ conn_num: int,
41
+ n_pre: int,
42
+ n_post: int,
43
+ seed: int | None,
44
+ method: str
45
+ ):
46
+ rng = random.default_rng(seed)
47
+
48
+ if method == 'vmap':
49
+ @transform.vmap(axis_size=n_pre)
50
+ def rand_indices():
51
+ return rng.choice(n_post, size=(conn_num,), replace=False)
52
+
53
+ return rand_indices()
54
+
55
+ elif method == 'for_loop':
56
+ return for_loop(
57
+ lambda *args: rng.choice(n_post, size=(conn_num,), replace=False),
58
+ length=n_pre
59
+ )
60
+
61
+ else:
62
+ raise ValueError(f"Unknown method: {method}")
63
+
64
+
65
+ class FixedNumConn(Module):
66
+ """
67
+ The ``FixedNumConn`` module implements a fixed probability connection with CSR sparse data structure.
68
+
69
+ Parameters
70
+ ----------
71
+ in_size : Size
72
+ Number of pre-synaptic neurons, i.e., input size.
73
+ out_size : Size
74
+ Number of post-synaptic neurons, i.e., output size.
75
+ conn_num : float, int
76
+ If it is a float, representing the probability of connection, i.e., connection probability.
77
+
78
+ If it is an integer, representing the number of connections.
79
+ conn_weight : float or callable or jax.Array or brainunit.Quantity
80
+ Maximum synaptic conductance, i.e., synaptic weight.
81
+ efferent_target : str, optional
82
+ The target of the connection. Default is 'post', meaning that each pre-synaptic neuron connects to
83
+ a fixed number of post-synaptic neurons. The connection number is determined by the value of ``n_conn``.
84
+
85
+ If 'pre', each post-synaptic neuron connects to a fixed number of pre-synaptic neurons.
86
+ conn_init : str, optional
87
+ The initialization method of the connection weight. Default is 'vmap', meaning that the connection weight
88
+ is initialized by parallelized across multiple threads.
89
+
90
+ If 'for_loop', the connection weight is initialized by a for loop.
91
+ allow_multi_conn : bool, optional
92
+ Whether multiple connections are allowed from a single pre-synaptic neuron.
93
+ Default is True, meaning that a value of ``a`` can be selected multiple times.
94
+ seed: int, optional
95
+ Random seed. Default is None. If None, the default random seed will be used.
96
+ name : str, optional
97
+ Name of the module.
98
+ """
99
+
100
+ __module__ = 'brainstate.nn'
101
+
102
+ def __init__(
103
+ self,
104
+ in_size: Size,
105
+ out_size: Size,
106
+ conn_num: Union[int, float],
107
+ conn_weight: Union[Callable, ArrayLike],
108
+ efferent_target: str = 'post', # 'pre' or 'post'
109
+ afferent_ratio: Union[int, float] = 1.,
110
+ allow_multi_conn: bool = True,
111
+ seed: Optional[int] = None,
112
+ name: Optional[str] = None,
113
+ conn_init: str = 'vmap', # 'vmap' or 'for_loop'
114
+ param_type: type = ParamState,
115
+ ):
116
+ super().__init__(name=name)
117
+
118
+ # network parameters
119
+ self.in_size = in_size
120
+ self.out_size = out_size
121
+ self.efferent_target = efferent_target
122
+ assert efferent_target in ('pre', 'post'), 'The target of the connection must be either "pre" or "post".'
123
+ assert 0. <= afferent_ratio <= 1., 'Afferent ratio must be in [0, 1].'
124
+ if isinstance(conn_num, float):
125
+ assert 0. <= conn_num <= 1., 'Connection probability must be in [0, 1].'
126
+ conn_num = (int(self.out_size[-1] * conn_num)
127
+ if efferent_target == 'post' else
128
+ int(self.in_size[-1] * conn_num))
129
+ assert isinstance(conn_num, int), 'Connection number must be an integer.'
130
+ self.conn_num = conn_num
131
+ self.seed = seed
132
+ self.allow_multi_conn = allow_multi_conn
133
+
134
+ # connections
135
+ if self.conn_num >= 1:
136
+ if self.efferent_target == 'post':
137
+ n_post = self.out_size[-1]
138
+ n_pre = self.in_size[-1]
139
+ else:
140
+ n_post = self.in_size[-1]
141
+ n_pre = self.out_size[-1]
142
+
143
+ with jax.ensure_compile_time_eval():
144
+ if allow_multi_conn:
145
+ rng = np.random if seed is None else np.random.RandomState(seed)
146
+ indices = rng.randint(0, n_post, size=(n_pre, self.conn_num))
147
+ else:
148
+ indices = init_indices_without_replace(self.conn_num, n_pre, n_post, seed, conn_init)
149
+ indices = u.math.asarray(indices, dtype=environ.ditype())
150
+
151
+ if afferent_ratio == 1.:
152
+ conn_weight = u.math.asarray(init.param(conn_weight, (n_pre, self.conn_num), allow_none=False))
153
+ self.weight = param_type(conn_weight)
154
+ csr = (
155
+ brainevent.FixedPostNumConn((conn_weight, indices), shape=(n_pre, n_post))
156
+ if self.efferent_target == 'post' else
157
+ brainevent.FixedPreNumConn((conn_weight, indices), shape=(n_pre, n_post))
158
+ )
159
+ self.conn = csr
160
+
161
+ else:
162
+ self.pre_selected = np.random.random(n_pre) < afferent_ratio
163
+ indices = indices[self.pre_selected].flatten()
164
+ conn_weight = u.math.asarray(init.param(conn_weight, (indices.size,), allow_none=False))
165
+ self.weight = param_type(conn_weight)
166
+ indptr = (jnp.arange(1, n_pre + 1) * self.conn_num -
167
+ jnp.cumsum(~self.pre_selected) * self.conn_num)
168
+ indptr = jnp.insert(indptr, 0, 0) # insert 0 at the beginning
169
+ csr = (
170
+ brainevent.CSR((conn_weight, indices, indptr), shape=(n_pre, n_post))
171
+ if self.efferent_target == 'post' else
172
+ brainevent.CSC((conn_weight, indices, indptr), shape=(n_pre, n_post))
173
+ )
174
+ self.conn = csr
175
+
176
+ else:
177
+ conn_weight = u.math.asarray(init.param(conn_weight, (), allow_none=False))
178
+ self.weight = FakeState(conn_weight)
179
+
180
+ def update(self, x) -> Union[jax.Array, u.Quantity]:
181
+ if self.conn_num >= 1:
182
+ csr = self.conn.with_data(self.weight.value)
183
+ return x @ csr
184
+ else:
185
+ weight = self.weight.value
186
+ r = u.math.zeros(x.shape[:-1] + (self.out_size[-1],), dtype=weight.dtype)
187
+ return u.maybe_decimal(u.Quantity(r, unit=u.get_unit(weight), dtype=environ.dftype()))
188
+
189
+
190
+ class EventFixedNumConn(FixedNumConn):
191
+ """
192
+ The FixedProb module implements a fixed probability connection with CSR sparse data structure.
193
+
194
+ Parameters
195
+ ----------
196
+ in_size : Size
197
+ Number of pre-synaptic neurons, i.e., input size.
198
+ out_size : Size
199
+ Number of post-synaptic neurons, i.e., output size.
200
+ conn_num : float, int
201
+ If it is a float, representing the probability of connection, i.e., connection probability.
202
+
203
+ If it is an integer, representing the number of connections.
204
+ conn_weight : float or callable or jax.Array or brainunit.Quantity
205
+ Maximum synaptic conductance, i.e., synaptic weight.
206
+ conn_target : str, optional
207
+ The target of the connection. Default is 'post', meaning that each pre-synaptic neuron connects to
208
+ a fixed number of post-synaptic neurons. The connection number is determined by the value of ``n_conn``.
209
+
210
+ If 'pre', each post-synaptic neuron connects to a fixed number of pre-synaptic neurons.
211
+ conn_init : str, optional
212
+ The initialization method of the connection weight. Default is 'vmap', meaning that the connection weight
213
+ is initialized by parallelized across multiple threads.
214
+
215
+ If 'for_loop', the connection weight is initialized by a for loop.
216
+ allow_multi_conn : bool, optional
217
+ Whether multiple connections are allowed from a single pre-synaptic neuron.
218
+ Default is True, meaning that a value of ``a`` can be selected multiple times.
219
+ seed: int, optional
220
+ Random seed. Default is None. If None, the default random seed will be used.
221
+ name : str, optional
222
+ Name of the module.
223
+ """
224
+
225
+ __module__ = 'brainstate.nn'
226
+
227
+ def update(self, spk: jax.Array) -> Union[jax.Array, u.Quantity]:
228
+ return super().update(
229
+ brainevent.EventArray(spk)
230
+ )
231
+
232
+
233
+ EventFixedProb = EventFixedNumConn