brainstate 0.1.10__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 (163) hide show
  1. brainstate/__init__.py +169 -58
  2. brainstate/_compatible_import.py +340 -148
  3. brainstate/_compatible_import_test.py +681 -0
  4. brainstate/_deprecation.py +210 -0
  5. brainstate/_deprecation_test.py +2319 -0
  6. brainstate/{util/error.py → _error.py} +45 -55
  7. brainstate/_state.py +1652 -1605
  8. brainstate/_state_test.py +52 -52
  9. brainstate/_utils.py +47 -47
  10. brainstate/environ.py +1495 -563
  11. brainstate/environ_test.py +1223 -62
  12. brainstate/graph/__init__.py +22 -29
  13. brainstate/graph/_node.py +240 -0
  14. brainstate/graph/_node_test.py +589 -0
  15. brainstate/graph/{_graph_operation.py → _operation.py} +1624 -1738
  16. brainstate/graph/_operation_test.py +1147 -0
  17. brainstate/mixin.py +1433 -365
  18. brainstate/mixin_test.py +1017 -77
  19. brainstate/nn/__init__.py +137 -135
  20. brainstate/nn/_activations.py +1100 -808
  21. brainstate/nn/_activations_test.py +354 -331
  22. brainstate/nn/_collective_ops.py +633 -514
  23. brainstate/nn/_collective_ops_test.py +774 -43
  24. brainstate/nn/_common.py +226 -178
  25. brainstate/nn/_common_test.py +154 -0
  26. brainstate/nn/_conv.py +2010 -501
  27. brainstate/nn/_conv_test.py +849 -238
  28. brainstate/nn/_delay.py +575 -588
  29. brainstate/nn/_delay_test.py +243 -238
  30. brainstate/nn/_dropout.py +618 -426
  31. brainstate/nn/_dropout_test.py +477 -100
  32. brainstate/nn/_dynamics.py +1267 -1343
  33. brainstate/nn/_dynamics_test.py +67 -78
  34. brainstate/nn/_elementwise.py +1298 -1119
  35. brainstate/nn/_elementwise_test.py +830 -169
  36. brainstate/nn/_embedding.py +408 -58
  37. brainstate/nn/_embedding_test.py +156 -0
  38. brainstate/nn/{_fixedprob.py → _event_fixedprob.py} +233 -239
  39. brainstate/nn/{_fixedprob_test.py → _event_fixedprob_test.py} +115 -114
  40. brainstate/nn/{_linear_mv.py → _event_linear.py} +83 -83
  41. brainstate/nn/{_linear_mv_test.py → _event_linear_test.py} +121 -120
  42. brainstate/nn/_exp_euler.py +254 -92
  43. brainstate/nn/_exp_euler_test.py +377 -35
  44. brainstate/nn/_linear.py +744 -424
  45. brainstate/nn/_linear_test.py +475 -107
  46. brainstate/nn/_metrics.py +1070 -0
  47. brainstate/nn/_metrics_test.py +611 -0
  48. brainstate/nn/_module.py +384 -377
  49. brainstate/nn/_module_test.py +40 -40
  50. brainstate/nn/_normalizations.py +1334 -975
  51. brainstate/nn/_normalizations_test.py +699 -73
  52. brainstate/nn/_paddings.py +1020 -0
  53. brainstate/nn/_paddings_test.py +723 -0
  54. brainstate/nn/_poolings.py +2239 -1177
  55. brainstate/nn/_poolings_test.py +953 -217
  56. brainstate/nn/{_rate_rnns.py → _rnns.py} +946 -554
  57. brainstate/nn/_rnns_test.py +593 -0
  58. brainstate/nn/_utils.py +216 -89
  59. brainstate/nn/_utils_test.py +402 -0
  60. brainstate/{init/_random_inits.py → nn/init.py} +809 -553
  61. brainstate/{init/_random_inits_test.py → nn/init_test.py} +180 -149
  62. brainstate/random/__init__.py +270 -24
  63. brainstate/random/_rand_funs.py +3938 -3616
  64. brainstate/random/_rand_funs_test.py +640 -567
  65. brainstate/random/_rand_seed.py +675 -210
  66. brainstate/random/_rand_seed_test.py +48 -48
  67. brainstate/random/_rand_state.py +1617 -1409
  68. brainstate/random/_rand_state_test.py +551 -0
  69. brainstate/transform/__init__.py +59 -0
  70. brainstate/transform/_ad_checkpoint.py +176 -0
  71. brainstate/{compile → transform}/_ad_checkpoint_test.py +49 -49
  72. brainstate/{augment → transform}/_autograd.py +1025 -778
  73. brainstate/{augment → transform}/_autograd_test.py +1289 -1289
  74. brainstate/transform/_conditions.py +316 -0
  75. brainstate/{compile → transform}/_conditions_test.py +220 -220
  76. brainstate/{compile → transform}/_error_if.py +94 -92
  77. brainstate/{compile → transform}/_error_if_test.py +52 -52
  78. brainstate/transform/_eval_shape.py +145 -0
  79. brainstate/{augment → transform}/_eval_shape_test.py +38 -38
  80. brainstate/{compile → transform}/_jit.py +399 -346
  81. brainstate/{compile → transform}/_jit_test.py +143 -143
  82. brainstate/{compile → transform}/_loop_collect_return.py +675 -536
  83. brainstate/{compile → transform}/_loop_collect_return_test.py +58 -58
  84. brainstate/{compile → transform}/_loop_no_collection.py +283 -184
  85. brainstate/{compile → transform}/_loop_no_collection_test.py +50 -50
  86. brainstate/transform/_make_jaxpr.py +2016 -0
  87. brainstate/transform/_make_jaxpr_test.py +1510 -0
  88. brainstate/transform/_mapping.py +529 -0
  89. brainstate/transform/_mapping_test.py +194 -0
  90. brainstate/{compile → transform}/_progress_bar.py +255 -202
  91. brainstate/{augment → transform}/_random.py +171 -151
  92. brainstate/{compile → transform}/_unvmap.py +256 -159
  93. brainstate/transform/_util.py +286 -0
  94. brainstate/typing.py +837 -304
  95. brainstate/typing_test.py +780 -0
  96. brainstate/util/__init__.py +27 -50
  97. brainstate/util/_others.py +1025 -0
  98. brainstate/util/_others_test.py +962 -0
  99. brainstate/util/_pretty_pytree.py +1301 -0
  100. brainstate/util/_pretty_pytree_test.py +675 -0
  101. brainstate/util/{pretty_repr.py → _pretty_repr.py} +462 -328
  102. brainstate/util/_pretty_repr_test.py +696 -0
  103. brainstate/util/filter.py +945 -469
  104. brainstate/util/filter_test.py +912 -0
  105. brainstate/util/struct.py +910 -523
  106. brainstate/util/struct_test.py +602 -0
  107. {brainstate-0.1.10.dist-info → brainstate-0.2.1.dist-info}/METADATA +108 -91
  108. brainstate-0.2.1.dist-info/RECORD +111 -0
  109. {brainstate-0.1.10.dist-info → brainstate-0.2.1.dist-info}/licenses/LICENSE +202 -202
  110. brainstate/augment/__init__.py +0 -30
  111. brainstate/augment/_eval_shape.py +0 -99
  112. brainstate/augment/_mapping.py +0 -1060
  113. brainstate/augment/_mapping_test.py +0 -597
  114. brainstate/compile/__init__.py +0 -38
  115. brainstate/compile/_ad_checkpoint.py +0 -204
  116. brainstate/compile/_conditions.py +0 -256
  117. brainstate/compile/_make_jaxpr.py +0 -888
  118. brainstate/compile/_make_jaxpr_test.py +0 -156
  119. brainstate/compile/_util.py +0 -147
  120. brainstate/functional/__init__.py +0 -27
  121. brainstate/graph/_graph_node.py +0 -244
  122. brainstate/graph/_graph_node_test.py +0 -73
  123. brainstate/graph/_graph_operation_test.py +0 -563
  124. brainstate/init/__init__.py +0 -26
  125. brainstate/init/_base.py +0 -52
  126. brainstate/init/_generic.py +0 -244
  127. brainstate/init/_regular_inits.py +0 -105
  128. brainstate/init/_regular_inits_test.py +0 -50
  129. brainstate/nn/_inputs.py +0 -608
  130. brainstate/nn/_ltp.py +0 -28
  131. brainstate/nn/_neuron.py +0 -705
  132. brainstate/nn/_neuron_test.py +0 -161
  133. brainstate/nn/_others.py +0 -46
  134. brainstate/nn/_projection.py +0 -486
  135. brainstate/nn/_rate_rnns_test.py +0 -63
  136. brainstate/nn/_readout.py +0 -209
  137. brainstate/nn/_readout_test.py +0 -53
  138. brainstate/nn/_stp.py +0 -236
  139. brainstate/nn/_synapse.py +0 -505
  140. brainstate/nn/_synapse_test.py +0 -131
  141. brainstate/nn/_synaptic_projection.py +0 -423
  142. brainstate/nn/_synouts.py +0 -162
  143. brainstate/nn/_synouts_test.py +0 -57
  144. brainstate/nn/metrics.py +0 -388
  145. brainstate/optim/__init__.py +0 -38
  146. brainstate/optim/_base.py +0 -64
  147. brainstate/optim/_lr_scheduler.py +0 -448
  148. brainstate/optim/_lr_scheduler_test.py +0 -50
  149. brainstate/optim/_optax_optimizer.py +0 -152
  150. brainstate/optim/_optax_optimizer_test.py +0 -53
  151. brainstate/optim/_sgd_optimizer.py +0 -1104
  152. brainstate/random/_random_for_unit.py +0 -52
  153. brainstate/surrogate.py +0 -1957
  154. brainstate/transform.py +0 -23
  155. brainstate/util/caller.py +0 -98
  156. brainstate/util/others.py +0 -540
  157. brainstate/util/pretty_pytree.py +0 -945
  158. brainstate/util/pretty_pytree_test.py +0 -159
  159. brainstate/util/pretty_table.py +0 -2954
  160. brainstate/util/scaling.py +0 -258
  161. brainstate-0.1.10.dist-info/RECORD +0 -130
  162. {brainstate-0.1.10.dist-info → brainstate-0.2.1.dist-info}/WHEEL +0 -0
  163. {brainstate-0.1.10.dist-info → brainstate-0.2.1.dist-info}/top_level.txt +0 -0
@@ -1,1119 +1,1298 @@
1
- # Copyright 2024 BDP 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
- # -*- coding: utf-8 -*-
17
-
18
- from typing import Optional
19
-
20
- import brainunit as u
21
- import jax.numpy as jnp
22
-
23
- from brainstate._state import ParamState
24
- from brainstate.typing import ArrayLike
25
- from . import _activations as F
26
- from ._module import ElementWiseBlock
27
-
28
- __all__ = [
29
- # activation functions
30
- 'Threshold', 'ReLU', 'RReLU', 'Hardtanh', 'ReLU6', 'Sigmoid', 'Hardsigmoid',
31
- 'Tanh', 'SiLU', 'Mish', 'Hardswish', 'ELU', 'CELU', 'SELU', 'GLU', 'GELU',
32
- 'Hardshrink', 'LeakyReLU', 'LogSigmoid', 'Softplus', 'Softshrink', 'PReLU',
33
- 'Softsign', 'Tanhshrink', 'Softmin', 'Softmax', 'Softmax2d', 'LogSoftmax',
34
-
35
- # others
36
- 'Identity', 'SpikeBitwise',
37
- ]
38
-
39
-
40
- class Threshold(ElementWiseBlock):
41
- r"""Thresholds each element of the input Tensor.
42
-
43
- Threshold is defined as:
44
-
45
- .. math::
46
- y =
47
- \begin{cases}
48
- x, &\text{ if } x > \text{threshold} \\
49
- \text{value}, &\text{ otherwise }
50
- \end{cases}
51
-
52
- Args:
53
- threshold: The value to threshold at
54
- value: The value to replace with
55
-
56
- Shape:
57
- - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
58
- - Output: :math:`(*)`, same shape as the input.
59
-
60
- Examples::
61
-
62
- >>> import brainstate.nn as nn
63
- >>> import brainstate
64
- >>> m = nn.Threshold(0.1, 20)
65
- >>> x = random.randn(2)
66
- >>> output = m(x)
67
- """
68
- __module__ = 'brainstate.nn'
69
- threshold: float
70
- value: float
71
-
72
- def __init__(self, threshold: float, value: float) -> None:
73
- super().__init__()
74
- self.threshold = threshold
75
- self.value = value
76
-
77
- def __call__(self, x: ArrayLike) -> ArrayLike:
78
- dtype = u.math.get_dtype(x)
79
- return jnp.where(x > jnp.asarray(self.threshold, dtype=dtype),
80
- x,
81
- jnp.asarray(self.value, dtype=dtype))
82
-
83
- def __repr__(self):
84
- return f'{self.__class__.__name__}(threshold={self.threshold}, value={self.value})'
85
-
86
-
87
- class ReLU(ElementWiseBlock):
88
- r"""Applies the rectified linear unit function element-wise:
89
-
90
- :math:`\text{ReLU}(x) = (x)^+ = \max(0, x)`
91
-
92
- Shape:
93
- - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
94
- - Output: :math:`(*)`, same shape as the input.
95
-
96
- Examples::
97
-
98
- >>> import brainstate.nn as nn
99
- >>> import brainstate as brainstate
100
- >>> m = nn.ReLU()
101
- >>> x = random.randn(2)
102
- >>> output = m(x)
103
-
104
-
105
- An implementation of CReLU - https://arxiv.org/abs/1603.05201
106
-
107
- >>> import brainstate.nn as nn
108
- >>> import brainstate as brainstate
109
- >>> m = nn.ReLU()
110
- >>> x = random.randn(2).unsqueeze(0)
111
- >>> output = jax.numpy.concat((m(x), m(-x)))
112
- """
113
- __module__ = 'brainstate.nn'
114
-
115
- def __call__(self, x: ArrayLike) -> ArrayLike:
116
- return F.relu(x)
117
-
118
- def __repr__(self):
119
- return f'{self.__class__.__name__}()'
120
-
121
-
122
- class RReLU(ElementWiseBlock):
123
- r"""Applies the randomized leaky rectified liner unit function, element-wise,
124
- as described in the paper:
125
-
126
- `Empirical Evaluation of Rectified Activations in Convolutional Network`_.
127
-
128
- The function is defined as:
129
-
130
- .. math::
131
- \text{RReLU}(x) =
132
- \begin{cases}
133
- x & \text{if } x \geq 0 \\
134
- ax & \text{ otherwise }
135
- \end{cases}
136
-
137
- where :math:`a` is randomly sampled from uniform distribution
138
- :math:`\mathcal{U}(\text{lower}, \text{upper})`.
139
-
140
- See: https://arxiv.org/pdf/1505.00853.pdf
141
-
142
- Args:
143
- lower: lower bound of the uniform distribution. Default: :math:`\frac{1}{8}`
144
- upper: upper bound of the uniform distribution. Default: :math:`\frac{1}{3}`
145
-
146
- Shape:
147
- - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
148
- - Output: :math:`(*)`, same shape as the input.
149
-
150
- Examples::
151
-
152
- >>> import brainstate.nn as nn
153
- >>> import brainstate as brainstate
154
- >>> m = nn.RReLU(0.1, 0.3)
155
- >>> x = random.randn(2)
156
- >>> output = m(x)
157
-
158
- .. _`Empirical Evaluation of Rectified Activations in Convolutional Network`:
159
- https://arxiv.org/abs/1505.00853
160
- """
161
- __module__ = 'brainstate.nn'
162
- lower: float
163
- upper: float
164
-
165
- def __init__(
166
- self,
167
- lower: float = 1. / 8,
168
- upper: float = 1. / 3,
169
- ):
170
- super().__init__()
171
- self.lower = lower
172
- self.upper = upper
173
-
174
- def __call__(self, x: ArrayLike) -> ArrayLike:
175
- return F.rrelu(x, self.lower, self.upper)
176
-
177
- def extra_repr(self):
178
- return f'{self.__class__.__name__}(lower={self.lower}, upper={self.upper})'
179
-
180
-
181
- class Hardtanh(ElementWiseBlock):
182
- r"""Applies the HardTanh function element-wise.
183
-
184
- HardTanh is defined as:
185
-
186
- .. math::
187
- \text{HardTanh}(x) = \begin{cases}
188
- \text{max\_val} & \text{ if } x > \text{ max\_val } \\
189
- \text{min\_val} & \text{ if } x < \text{ min\_val } \\
190
- x & \text{ otherwise } \\
191
- \end{cases}
192
-
193
- Args:
194
- min_val: minimum value of the linear region range. Default: -1
195
- max_val: maximum value of the linear region range. Default: 1
196
-
197
- Keyword arguments :attr:`min_value` and :attr:`max_value`
198
- have been deprecated in favor of :attr:`min_val` and :attr:`max_val`.
199
-
200
- Shape:
201
- - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
202
- - Output: :math:`(*)`, same shape as the input.
203
-
204
- Examples::
205
-
206
- >>> import brainstate.nn as nn
207
- >>> import brainstate as brainstate
208
- >>> m = nn.Hardtanh(-2, 2)
209
- >>> x = random.randn(2)
210
- >>> output = m(x)
211
- """
212
- __module__ = 'brainstate.nn'
213
- min_val: float
214
- max_val: float
215
-
216
- def __init__(
217
- self,
218
- min_val: float = -1.,
219
- max_val: float = 1.,
220
- ) -> None:
221
- super().__init__()
222
- self.min_val = min_val
223
- self.max_val = max_val
224
- assert self.max_val > self.min_val
225
-
226
- def __call__(self, x: ArrayLike) -> ArrayLike:
227
- return F.hard_tanh(x, self.min_val, self.max_val)
228
-
229
- def extra_repr(self) -> str:
230
- return f'{self.__class__.__name__}(min_val={self.min_val}, max_val={self.max_val})'
231
-
232
-
233
- class ReLU6(Hardtanh, ElementWiseBlock):
234
- r"""Applies the element-wise function:
235
-
236
- .. math::
237
- \text{ReLU6}(x) = \min(\max(0,x), 6)
238
-
239
- Shape:
240
- - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
241
- - Output: :math:`(*)`, same shape as the input.
242
-
243
- Examples::
244
-
245
- >>> import brainstate.nn as nn
246
- >>> import brainstate as brainstate
247
- >>> m = nn.ReLU6()
248
- >>> x = random.randn(2)
249
- >>> output = m(x)
250
- """
251
- __module__ = 'brainstate.nn'
252
-
253
- def __init__(self):
254
- super().__init__(0., 6.)
255
-
256
-
257
- class Sigmoid(ElementWiseBlock):
258
- r"""Applies the element-wise function:
259
-
260
- .. math::
261
- \text{Sigmoid}(x) = \sigma(x) = \frac{1}{1 + \exp(-x)}
262
-
263
-
264
- Shape:
265
- - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
266
- - Output: :math:`(*)`, same shape as the input.
267
-
268
- Examples::
269
-
270
- >>> import brainstate.nn as nn
271
- >>> import brainstate as brainstate
272
- >>> m = nn.Sigmoid()
273
- >>> x = random.randn(2)
274
- >>> output = m(x)
275
- """
276
- __module__ = 'brainstate.nn'
277
-
278
- def __call__(self, x: ArrayLike) -> ArrayLike:
279
- return F.sigmoid(x)
280
-
281
-
282
- class Hardsigmoid(ElementWiseBlock):
283
- r"""Applies the Hardsigmoid function element-wise.
284
-
285
- Hardsigmoid is defined as:
286
-
287
- .. math::
288
- \text{Hardsigmoid}(x) = \begin{cases}
289
- 0 & \text{if~} x \le -3, \\
290
- 1 & \text{if~} x \ge +3, \\
291
- x / 6 + 1 / 2 & \text{otherwise}
292
- \end{cases}
293
-
294
- Shape:
295
- - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
296
- - Output: :math:`(*)`, same shape as the input.
297
-
298
- Examples::
299
-
300
- >>> import brainstate.nn as nn
301
- >>> import brainstate as brainstate
302
- >>> m = nn.Hardsigmoid()
303
- >>> x = random.randn(2)
304
- >>> output = m(x)
305
- """
306
- __module__ = 'brainstate.nn'
307
-
308
- def __call__(self, x: ArrayLike) -> ArrayLike:
309
- return F.hard_sigmoid(x)
310
-
311
-
312
- class Tanh(ElementWiseBlock):
313
- r"""Applies the Hyperbolic Tangent (Tanh) function element-wise.
314
-
315
- Tanh is defined as:
316
-
317
- .. math::
318
- \text{Tanh}(x) = \tanh(x) = \frac{\exp(x) - \exp(-x)} {\exp(x) + \exp(-x)}
319
-
320
- Shape:
321
- - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
322
- - Output: :math:`(*)`, same shape as the input.
323
-
324
- Examples::
325
-
326
- >>> import brainstate.nn as nn
327
- >>> import brainstate as brainstate
328
- >>> m = nn.Tanh()
329
- >>> x = random.randn(2)
330
- >>> output = m(x)
331
- """
332
- __module__ = 'brainstate.nn'
333
-
334
- def __call__(self, x: ArrayLike) -> ArrayLike:
335
- return F.tanh(x)
336
-
337
-
338
- class SiLU(ElementWiseBlock):
339
- r"""Applies the Sigmoid Linear Unit (SiLU) function, element-wise.
340
- The SiLU function is also known as the swish function.
341
-
342
- .. math::
343
- \text{silu}(x) = x * \sigma(x), \text{where } \sigma(x) \text{ is the logistic sigmoid.}
344
-
345
- .. note::
346
- See `Gaussian Error Linear Units (GELUs) <https://arxiv.org/abs/1606.08415>`_
347
- where the SiLU (Sigmoid Linear Unit) was originally coined, and see
348
- `Sigmoid-Weighted Linear Units for Neural Network Function Approximation
349
- in Reinforcement Learning <https://arxiv.org/abs/1702.03118>`_ and `Swish:
350
- a Self-Gated Activation Function <https://arxiv.org/abs/1710.05941v1>`_
351
- where the SiLU was experimented with later.
352
-
353
- Shape:
354
- - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
355
- - Output: :math:`(*)`, same shape as the input.
356
-
357
- Examples::
358
-
359
- >>> import brainstate.nn as nn
360
- >>> m = nn.SiLU()
361
- >>> x = random.randn(2)
362
- >>> output = m(x)
363
- """
364
- __module__ = 'brainstate.nn'
365
-
366
- def __call__(self, x: ArrayLike) -> ArrayLike:
367
- return F.silu(x)
368
-
369
-
370
- class Mish(ElementWiseBlock):
371
- r"""Applies the Mish function, element-wise.
372
- Mish: A Self Regularized Non-Monotonic Neural Activation Function.
373
-
374
- .. math::
375
- \text{Mish}(x) = x * \text{Tanh}(\text{Softplus}(x))
376
-
377
- .. note::
378
- See `Mish: A Self Regularized Non-Monotonic Neural Activation Function <https://arxiv.org/abs/1908.08681>`_
379
-
380
-
381
- Shape:
382
- - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
383
- - Output: :math:`(*)`, same shape as the input.
384
-
385
- Examples::
386
-
387
- >>> import brainstate.nn as nn
388
- >>> import brainstate as brainstate
389
- >>> m = nn.Mish()
390
- >>> x = random.randn(2)
391
- >>> output = m(x)
392
- """
393
- __module__ = 'brainstate.nn'
394
-
395
- def __call__(self, x: ArrayLike) -> ArrayLike:
396
- return F.mish(x)
397
-
398
-
399
- class Hardswish(ElementWiseBlock):
400
- r"""Applies the Hardswish function, element-wise, as described in the paper:
401
- `Searching for MobileNetV3 <https://arxiv.org/abs/1905.02244>`_.
402
-
403
- Hardswish is defined as:
404
-
405
- .. math::
406
- \text{Hardswish}(x) = \begin{cases}
407
- 0 & \text{if~} x \le -3, \\
408
- x & \text{if~} x \ge +3, \\
409
- x \cdot (x + 3) /6 & \text{otherwise}
410
- \end{cases}
411
-
412
-
413
- Shape:
414
- - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
415
- - Output: :math:`(*)`, same shape as the input.
416
-
417
- Examples::
418
-
419
- >>> import brainstate.nn as nn
420
- >>> import brainstate as brainstate
421
- >>> m = nn.Hardswish()
422
- >>> x = random.randn(2)
423
- >>> output = m(x)
424
- """
425
- __module__ = 'brainstate.nn'
426
-
427
- def __call__(self, x: ArrayLike) -> ArrayLike:
428
- return F.hard_swish(x)
429
-
430
-
431
- class ELU(ElementWiseBlock):
432
- r"""Applies the Exponential Linear Unit (ELU) function, element-wise, as described
433
- in the paper: `Fast and Accurate Deep Network Learning by Exponential Linear
434
- Units (ELUs) <https://arxiv.org/abs/1511.07289>`__.
435
-
436
- ELU is defined as:
437
-
438
- .. math::
439
- \text{ELU}(x) = \begin{cases}
440
- x, & \text{ if } x > 0\\
441
- \alpha * (\exp(x) - 1), & \text{ if } x \leq 0
442
- \end{cases}
443
-
444
- Args:
445
- alpha: the :math:`\alpha` value for the ELU formulation. Default: 1.0
446
-
447
- Shape:
448
- - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
449
- - Output: :math:`(*)`, same shape as the input.
450
-
451
- Examples::
452
-
453
- >>> import brainstate.nn as nn
454
- >>> import brainstate as brainstate
455
- >>> m = nn.ELU()
456
- >>> x = random.randn(2)
457
- >>> output = m(x)
458
- """
459
- __module__ = 'brainstate.nn'
460
- alpha: float
461
-
462
- def __init__(self, alpha: float = 1.) -> None:
463
- super().__init__()
464
- self.alpha = alpha
465
-
466
- def __call__(self, x: ArrayLike) -> ArrayLike:
467
- return F.elu(x, self.alpha)
468
-
469
- def extra_repr(self) -> str:
470
- return f'{self.__class__.__name__}(alpha={self.alpha})'
471
-
472
-
473
- class CELU(ElementWiseBlock):
474
- r"""Applies the element-wise function:
475
-
476
- .. math::
477
- \text{CELU}(x) = \max(0,x) + \min(0, \alpha * (\exp(x/\alpha) - 1))
478
-
479
- More details can be found in the paper `Continuously Differentiable Exponential Linear Units`_ .
480
-
481
- Args:
482
- alpha: the :math:`\alpha` value for the CELU formulation. Default: 1.0
483
-
484
- Shape:
485
- - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
486
- - Output: :math:`(*)`, same shape as the input.
487
-
488
- Examples::
489
-
490
- >>> import brainstate.nn as nn
491
- >>> import brainstate as brainstate
492
- >>> m = nn.CELU()
493
- >>> x = random.randn(2)
494
- >>> output = m(x)
495
-
496
- .. _`Continuously Differentiable Exponential Linear Units`:
497
- https://arxiv.org/abs/1704.07483
498
- """
499
- __module__ = 'brainstate.nn'
500
- alpha: float
501
-
502
- def __init__(self, alpha: float = 1.) -> None:
503
- super().__init__()
504
- self.alpha = alpha
505
-
506
- def __call__(self, x: ArrayLike) -> ArrayLike:
507
- return F.celu(x, self.alpha)
508
-
509
- def extra_repr(self) -> str:
510
- return f'{self.__class__.__name__}(alpha={self.alpha})'
511
-
512
-
513
- class SELU(ElementWiseBlock):
514
- r"""Applied element-wise, as:
515
-
516
- .. math::
517
- \text{SELU}(x) = \text{scale} * (\max(0,x) + \min(0, \alpha * (\exp(x) - 1)))
518
-
519
- with :math:`\alpha = 1.6732632423543772848170429916717` and
520
- :math:`\text{scale} = 1.0507009873554804934193349852946`.
521
-
522
- More details can be found in the paper `Self-Normalizing Neural Networks`_ .
523
-
524
-
525
- Shape:
526
- - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
527
- - Output: :math:`(*)`, same shape as the input.
528
-
529
- Examples::
530
-
531
- >>> import brainstate.nn as nn
532
- >>> import brainstate as brainstate
533
- >>> m = nn.SELU()
534
- >>> x = random.randn(2)
535
- >>> output = m(x)
536
-
537
- .. _Self-Normalizing Neural Networks: https://arxiv.org/abs/1706.02515
538
- """
539
- __module__ = 'brainstate.nn'
540
-
541
- def __call__(self, x: ArrayLike) -> ArrayLike:
542
- return F.selu(x)
543
-
544
-
545
- class GLU(ElementWiseBlock):
546
- r"""Applies the gated linear unit function
547
- :math:`{GLU}(a, b)= a \otimes \sigma(b)` where :math:`a` is the first half
548
- of the input matrices and :math:`b` is the second half.
549
-
550
- Args:
551
- dim (int): the dimension on which to split the input. Default: -1
552
-
553
- Shape:
554
- - Input: :math:`(\ast_1, N, \ast_2)` where `*` means, any number of additional
555
- dimensions
556
- - Output: :math:`(\ast_1, M, \ast_2)` where :math:`M=N/2`
557
-
558
- Examples::
559
-
560
- >>> import brainstate.nn as nn
561
- >>> import brainstate as brainstate
562
- >>> m = nn.GLU()
563
- >>> x = random.randn(4, 2)
564
- >>> output = m(x)
565
- """
566
- __module__ = 'brainstate.nn'
567
- dim: int
568
-
569
- def __init__(self, dim: int = -1) -> None:
570
- super().__init__()
571
- self.dim = dim
572
-
573
- def __call__(self, x: ArrayLike) -> ArrayLike:
574
- return F.glu(x, self.dim)
575
-
576
- def __repr__(self):
577
- return f'{self.__class__.__name__}(dim={self.dim})'
578
-
579
-
580
- class GELU(ElementWiseBlock):
581
- r"""Applies the Gaussian Error Linear Units function:
582
-
583
- .. math:: \text{GELU}(x) = x * \Phi(x)
584
-
585
- where :math:`\Phi(x)` is the Cumulative Distribution Function for Gaussian Distribution.
586
-
587
- When the approximate argument is 'tanh', Gelu is estimated with:
588
-
589
- .. math:: \text{GELU}(x) = 0.5 * x * (1 + \text{Tanh}(\sqrt(2 / \pi) * (x + 0.044715 * x^3)))
590
-
591
- Args:
592
- approximate (str, optional): the gelu approximation algorithm to use:
593
- ``'none'`` | ``'tanh'``. Default: ``'none'``
594
-
595
- Shape:
596
- - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
597
- - Output: :math:`(*)`, same shape as the input.
598
-
599
- Examples::
600
-
601
- >>> import brainstate.nn as nn
602
- >>> import brainstate as brainstate
603
- >>> m = nn.GELU()
604
- >>> x = random.randn(2)
605
- >>> output = m(x)
606
- """
607
- __module__ = 'brainstate.nn'
608
- approximate: bool
609
-
610
- def __init__(self, approximate: bool = False) -> None:
611
- super().__init__()
612
- self.approximate = approximate
613
-
614
- def __call__(self, x: ArrayLike) -> ArrayLike:
615
- return F.gelu(x, approximate=self.approximate)
616
-
617
- def __repr__(self):
618
- return f'{self.__class__.__name__}(approximate={self.approximate})'
619
-
620
-
621
- class Hardshrink(ElementWiseBlock):
622
- r"""Applies the Hard Shrinkage (Hardshrink) function element-wise.
623
-
624
- Hardshrink is defined as:
625
-
626
- .. math::
627
- \text{HardShrink}(x) =
628
- \begin{cases}
629
- x, & \text{ if } x > \lambda \\
630
- x, & \text{ if } x < -\lambda \\
631
- 0, & \text{ otherwise }
632
- \end{cases}
633
-
634
- Args:
635
- lambd: the :math:`\lambda` value for the Hardshrink formulation. Default: 0.5
636
-
637
- Shape:
638
- - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
639
- - Output: :math:`(*)`, same shape as the input.
640
-
641
- Examples::
642
-
643
- >>> import brainstate.nn as nn
644
- >>> import brainstate as brainstate
645
- >>> m = nn.Hardshrink()
646
- >>> x = random.randn(2)
647
- >>> output = m(x)
648
- """
649
- __module__ = 'brainstate.nn'
650
- lambd: float
651
-
652
- def __init__(self, lambd: float = 0.5) -> None:
653
- super().__init__()
654
- self.lambd = lambd
655
-
656
- def __call__(self, x: ArrayLike) -> ArrayLike:
657
- return F.hard_shrink(x, self.lambd)
658
-
659
- def __repr__(self):
660
- return f'{self.__class__.__name__}(lambd={self.lambd})'
661
-
662
-
663
- class LeakyReLU(ElementWiseBlock):
664
- r"""Applies the element-wise function:
665
-
666
- .. math::
667
- \text{LeakyReLU}(x) = \max(0, x) + \text{negative\_slope} * \min(0, x)
668
-
669
-
670
- or
671
-
672
- .. math::
673
- \text{LeakyReLU}(x) =
674
- \begin{cases}
675
- x, & \text{ if } x \geq 0 \\
676
- \text{negative\_slope} \times x, & \text{ otherwise }
677
- \end{cases}
678
-
679
- Args:
680
- negative_slope: Controls the angle of the negative slope (which is used for
681
- negative input values). Default: 1e-2
682
-
683
- Shape:
684
- - Input: :math:`(*)` where `*` means, any number of additional
685
- dimensions
686
- - Output: :math:`(*)`, same shape as the input
687
-
688
- Examples::
689
-
690
- >>> import brainstate.nn as nn
691
- >>> import brainstate as brainstate
692
- >>> m = nn.LeakyReLU(0.1)
693
- >>> x = random.randn(2)
694
- >>> output = m(x)
695
- """
696
- __module__ = 'brainstate.nn'
697
- negative_slope: float
698
-
699
- def __init__(self, negative_slope: float = 1e-2) -> None:
700
- super().__init__()
701
- self.negative_slope = negative_slope
702
-
703
- def __call__(self, x: ArrayLike) -> ArrayLike:
704
- return F.leaky_relu(x, self.negative_slope)
705
-
706
- def __repr__(self):
707
- return f'{self.__class__.__name__}(negative_slope={self.negative_slope})'
708
-
709
-
710
- class LogSigmoid(ElementWiseBlock):
711
- r"""Applies the element-wise function:
712
-
713
- .. math::
714
- \text{LogSigmoid}(x) = \log\left(\frac{ 1 }{ 1 + \exp(-x)}\right)
715
-
716
- Shape:
717
- - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
718
- - Output: :math:`(*)`, same shape as the input.
719
-
720
- Examples::
721
-
722
- >>> import brainstate.nn as nn
723
- >>> import brainstate as brainstate
724
- >>> m = nn.LogSigmoid()
725
- >>> x = random.randn(2)
726
- >>> output = m(x)
727
- """
728
- __module__ = 'brainstate.nn'
729
-
730
- def __call__(self, x: ArrayLike) -> ArrayLike:
731
- return F.log_sigmoid(x)
732
-
733
-
734
- class Softplus(ElementWiseBlock):
735
- r"""Applies the Softplus function :math:`\text{Softplus}(x) = \frac{1}{\beta} *
736
- \log(1 + \exp(\beta * x))` element-wise.
737
-
738
- SoftPlus is a smooth approximation to the ReLU function and can be used
739
- to constrain the output of a machine to always be positive.
740
-
741
- For numerical stability the implementation reverts to the linear function
742
- when :math:`input \times \beta > threshold`.
743
-
744
- Shape:
745
- - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
746
- - Output: :math:`(*)`, same shape as the input.
747
-
748
- Examples::
749
-
750
- >>> import brainstate.nn as nn
751
- >>> import brainstate as brainstate
752
- >>> m = nn.Softplus()
753
- >>> x = random.randn(2)
754
- >>> output = m(x)
755
- """
756
- __module__ = 'brainstate.nn'
757
-
758
- def __call__(self, x: ArrayLike) -> ArrayLike:
759
- return F.softplus(x)
760
-
761
-
762
- class Softshrink(ElementWiseBlock):
763
- r"""Applies the soft shrinkage function elementwise:
764
-
765
- .. math::
766
- \text{SoftShrinkage}(x) =
767
- \begin{cases}
768
- x - \lambda, & \text{ if } x > \lambda \\
769
- x + \lambda, & \text{ if } x < -\lambda \\
770
- 0, & \text{ otherwise }
771
- \end{cases}
772
-
773
- Args:
774
- lambd: the :math:`\lambda` (must be no less than zero) value for the Softshrink formulation. Default: 0.5
775
-
776
- Shape:
777
- - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
778
- - Output: :math:`(*)`, same shape as the input.
779
-
780
- Examples::
781
-
782
- >>> import brainstate.nn as nn
783
- >>> import brainstate as brainstate
784
- >>> m = nn.Softshrink()
785
- >>> x = random.randn(2)
786
- >>> output = m(x)
787
- """
788
- __module__ = 'brainstate.nn'
789
- lambd: float
790
-
791
- def __init__(self, lambd: float = 0.5) -> None:
792
- super().__init__()
793
- self.lambd = lambd
794
-
795
- def __call__(self, x: ArrayLike) -> ArrayLike:
796
- return F.soft_shrink(x, self.lambd)
797
-
798
- def __repr__(self):
799
- return f'{self.__class__.__name__}(lambd={self.lambd})'
800
-
801
-
802
- class PReLU(ElementWiseBlock):
803
- r"""Applies the element-wise function:
804
-
805
- .. math::
806
- \text{PReLU}(x) = \max(0,x) + a * \min(0,x)
807
-
808
- or
809
-
810
- .. math::
811
- \text{PReLU}(x) =
812
- \begin{cases}
813
- x, & \text{ if } x \geq 0 \\
814
- ax, & \text{ otherwise }
815
- \end{cases}
816
-
817
- Here :math:`a` is a learnable parameter. When called without arguments, `nn.PReLU()` uses a single
818
- parameter :math:`a` across all input channels. If called with `nn.PReLU(nChannels)`,
819
- a separate :math:`a` is used for each input channel.
820
-
821
-
822
- .. note::
823
- weight decay should not be used when learning :math:`a` for good performance.
824
-
825
- .. note::
826
- Channel dim is the 2nd dim of input. When input has dims < 2, then there is
827
- no channel dim and the number of channels = 1.
828
-
829
- Args:
830
- num_parameters (int): number of :math:`a` to learn.
831
- Although it takes an int as input, there is only two values are legitimate:
832
- 1, or the number of channels at input. Default: 1
833
- init (float): the initial value of :math:`a`. Default: 0.25
834
-
835
- Shape:
836
- - Input: :math:`( *)` where `*` means, any number of additional
837
- dimensions.
838
- - Output: :math:`(*)`, same shape as the input.
839
-
840
- Attributes:
841
- weight (Tensor): the learnable weights of shape (:attr:`num_parameters`).
842
-
843
- Examples::
844
-
845
- >>> import brainstate as brainstate
846
- >>> m = brainstate.nn.PReLU()
847
- >>> x = brainstate.random.randn(2)
848
- >>> output = m(x)
849
- """
850
- __module__ = 'brainstate.nn'
851
- num_parameters: int
852
-
853
- def __init__(self, num_parameters: int = 1, init: float = 0.25, dtype=None) -> None:
854
- super().__init__()
855
- self.num_parameters = num_parameters
856
- self.weight = ParamState(jnp.ones(num_parameters, dtype=dtype) * init)
857
-
858
- def __call__(self, x: ArrayLike) -> ArrayLike:
859
- return F.prelu(x, self.weight.value)
860
-
861
- def __repr__(self):
862
- return f'{self.__class__.__name__}(num_parameters={self.num_parameters})'
863
-
864
-
865
- class Softsign(ElementWiseBlock):
866
- r"""Applies the element-wise function:
867
-
868
- .. math::
869
- \text{SoftSign}(x) = \frac{x}{ 1 + |x|}
870
-
871
- Shape:
872
- - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
873
- - Output: :math:`(*)`, same shape as the input.
874
-
875
- Examples::
876
-
877
- >>> import brainstate.nn as nn
878
- >>> import brainstate as brainstate
879
- >>> m = nn.Softsign()
880
- >>> x = random.randn(2)
881
- >>> output = m(x)
882
- """
883
- __module__ = 'brainstate.nn'
884
-
885
- def __call__(self, x: ArrayLike) -> ArrayLike:
886
- return F.soft_sign(x)
887
-
888
-
889
- class Tanhshrink(ElementWiseBlock):
890
- r"""Applies the element-wise function:
891
-
892
- .. math::
893
- \text{Tanhshrink}(x) = x - \tanh(x)
894
-
895
- Shape:
896
- - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
897
- - Output: :math:`(*)`, same shape as the input.
898
-
899
- Examples::
900
-
901
- >>> import brainstate.nn as nn
902
- >>> import brainstate as brainstate
903
- >>> m = nn.Tanhshrink()
904
- >>> x = random.randn(2)
905
- >>> output = m(x)
906
- """
907
- __module__ = 'brainstate.nn'
908
-
909
- def __call__(self, x: ArrayLike) -> ArrayLike:
910
- return F.tanh_shrink(x)
911
-
912
-
913
- class Softmin(ElementWiseBlock):
914
- r"""Applies the Softmin function to an n-dimensional input Tensor
915
- rescaling them so that the elements of the n-dimensional output Tensor
916
- lie in the range `[0, 1]` and sum to 1.
917
-
918
- Softmin is defined as:
919
-
920
- .. math::
921
- \text{Softmin}(x_{i}) = \frac{\exp(-x_i)}{\sum_j \exp(-x_j)}
922
-
923
- Shape:
924
- - Input: :math:`(*)` where `*` means, any number of additional
925
- dimensions
926
- - Output: :math:`(*)`, same shape as the input
927
-
928
- Args:
929
- dim (int): A dimension along which Softmin will be computed (so every slice
930
- along dim will sum to 1).
931
-
932
- Returns:
933
- a Tensor of the same dimension and shape as the input, with
934
- values in the range [0, 1]
935
-
936
- Examples::
937
-
938
- >>> import brainstate.nn as nn
939
- >>> import brainstate as brainstate
940
- >>> m = nn.Softmin(dim=1)
941
- >>> x = random.randn(2, 3)
942
- >>> output = m(x)
943
- """
944
- __module__ = 'brainstate.nn'
945
- dim: Optional[int]
946
-
947
- def __init__(self, dim: Optional[int] = None) -> None:
948
- super().__init__()
949
- self.dim = dim
950
-
951
- def __call__(self, x: ArrayLike) -> ArrayLike:
952
- return F.softmin(x, self.dim)
953
-
954
- def __repr__(self):
955
- return f'{self.__class__.__name__}(dim={self.dim})'
956
-
957
-
958
- class Softmax(ElementWiseBlock):
959
- r"""Applies the Softmax function to an n-dimensional input Tensor
960
- rescaling them so that the elements of the n-dimensional output Tensor
961
- lie in the range [0,1] and sum to 1.
962
-
963
- Softmax is defined as:
964
-
965
- .. math::
966
- \text{Softmax}(x_{i}) = \frac{\exp(x_i)}{\sum_j \exp(x_j)}
967
-
968
- When the input Tensor is a sparse tensor then the unspecified
969
- values are treated as ``-inf``.
970
-
971
- Shape:
972
- - Input: :math:`(*)` where `*` means, any number of additional
973
- dimensions
974
- - Output: :math:`(*)`, same shape as the input
975
-
976
- Returns:
977
- a Tensor of the same dimension and shape as the input with
978
- values in the range [0, 1]
979
-
980
- Args:
981
- dim (int): A dimension along which Softmax will be computed (so every slice
982
- along dim will sum to 1).
983
-
984
- .. note::
985
- This module doesn't work directly with NLLLoss,
986
- which expects the Log to be computed between the Softmax and itself.
987
- Use `LogSoftmax` instead (it's faster and has better numerical properties).
988
-
989
- Examples::
990
-
991
- >>> import brainstate.nn as nn
992
- >>> import brainstate as brainstate
993
- >>> m = nn.Softmax(dim=1)
994
- >>> x = random.randn(2, 3)
995
- >>> output = m(x)
996
-
997
- """
998
- __module__ = 'brainstate.nn'
999
- dim: Optional[int]
1000
-
1001
- def __init__(self, dim: Optional[int] = None) -> None:
1002
- super().__init__()
1003
- self.dim = dim
1004
-
1005
- def __call__(self, x: ArrayLike) -> ArrayLike:
1006
- return F.softmax(x, self.dim)
1007
-
1008
- def __repr__(self):
1009
- return f'{self.__class__.__name__}(dim={self.dim})'
1010
-
1011
-
1012
- class Softmax2d(ElementWiseBlock):
1013
- r"""Applies SoftMax over features to each spatial location.
1014
-
1015
- When given an image of ``Channels x Height x Width``, it will
1016
- apply `Softmax` to each location :math:`(Channels, h_i, w_j)`
1017
-
1018
- Shape:
1019
- - Input: :math:`(N, C, H, W)` or :math:`(C, H, W)`.
1020
- - Output: :math:`(N, C, H, W)` or :math:`(C, H, W)` (same shape as input)
1021
-
1022
- Returns:
1023
- a Tensor of the same dimension and shape as the input with
1024
- values in the range [0, 1]
1025
-
1026
- Examples::
1027
-
1028
- >>> import brainstate.nn as nn
1029
- >>> import brainstate as brainstate
1030
- >>> m = nn.Softmax2d()
1031
- >>> # you softmax over the 2nd dimension
1032
- >>> x = random.randn(2, 3, 12, 13)
1033
- >>> output = m(x)
1034
- """
1035
- __module__ = 'brainstate.nn'
1036
-
1037
- def __call__(self, x: ArrayLike) -> ArrayLike:
1038
- assert x.ndim == 4 or x.ndim == 3, 'Softmax2d requires a 3D or 4D tensor as input'
1039
- return F.softmax(x, -3)
1040
-
1041
-
1042
- class LogSoftmax(ElementWiseBlock):
1043
- r"""Applies the :math:`\log(\text{Softmax}(x))` function to an n-dimensional
1044
- input Tensor. The LogSoftmax formulation can be simplified as:
1045
-
1046
- .. math::
1047
- \text{LogSoftmax}(x_{i}) = \log\left(\frac{\exp(x_i) }{ \sum_j \exp(x_j)} \right)
1048
-
1049
- Shape:
1050
- - Input: :math:`(*)` where `*` means, any number of additional
1051
- dimensions
1052
- - Output: :math:`(*)`, same shape as the input
1053
-
1054
- Args:
1055
- dim (int): A dimension along which LogSoftmax will be computed.
1056
-
1057
- Returns:
1058
- a Tensor of the same dimension and shape as the input with
1059
- values in the range [-inf, 0)
1060
-
1061
- Examples::
1062
-
1063
- >>> import brainstate.nn as nn
1064
- >>> import brainstate as brainstate
1065
- >>> m = nn.LogSoftmax(dim=1)
1066
- >>> x = random.randn(2, 3)
1067
- >>> output = m(x)
1068
- """
1069
- __module__ = 'brainstate.nn'
1070
- dim: Optional[int]
1071
-
1072
- def __init__(self, dim: Optional[int] = None) -> None:
1073
- super().__init__()
1074
- self.dim = dim
1075
-
1076
- def __call__(self, x: ArrayLike) -> ArrayLike:
1077
- return F.log_softmax(x, self.dim)
1078
-
1079
- def __repr__(self):
1080
- return f'{self.__class__.__name__}(dim={self.dim})'
1081
-
1082
-
1083
- class Identity(ElementWiseBlock):
1084
- r"""A placeholder identity operator that is argument-insensitive.
1085
- """
1086
- __module__ = 'brainstate.nn'
1087
-
1088
- def __call__(self, x):
1089
- return x
1090
-
1091
-
1092
- class SpikeBitwise(ElementWiseBlock):
1093
- r"""Bitwise addition for the spiking inputs.
1094
-
1095
- .. math::
1096
-
1097
- \begin{array}{ccc}
1098
- \hline \text { Mode } & \text { Expression for } \mathrm{g}(\mathrm{x}, \mathrm{y}) & \text { Code for } \mathrm{g}(\mathrm{x}, \mathrm{y}) \\
1099
- \hline \text { ADD } & x+y & x+y \\
1100
- \text { AND } & x \cap y & x \cdot y \\
1101
- \text { IAND } & (\neg x) \cap y & (1-x) \cdot y \\
1102
- \text { OR } & x \cup y & (x+y)-(x \cdot y) \\
1103
- \hline
1104
- \end{array}
1105
-
1106
- Args:
1107
- op: str. The bitwise operation.
1108
- name: str. The name of the dynamic system.
1109
- """
1110
- __module__ = 'brainstate.nn'
1111
-
1112
- def __init__(self,
1113
- op: str = 'add',
1114
- name: Optional[str] = None) -> None:
1115
- super().__init__(name=name)
1116
- self.op = op
1117
-
1118
- def __call__(self, x, y):
1119
- return F.spike_bitwise(x, y, self.op)
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
+ # -*- coding: utf-8 -*-
17
+
18
+ from typing import Optional
19
+
20
+ import brainunit as u
21
+ import jax.numpy as jnp
22
+
23
+ from brainstate._state import ParamState
24
+ from brainstate.typing import ArrayLike
25
+ from . import _activations as F
26
+ from ._module import ElementWiseBlock
27
+
28
+ __all__ = [
29
+ # activation functions
30
+ 'Threshold', 'ReLU', 'RReLU', 'Hardtanh', 'ReLU6', 'Sigmoid', 'Hardsigmoid',
31
+ 'Tanh', 'SiLU', 'Mish', 'Hardswish', 'ELU', 'CELU', 'SELU', 'GLU', 'GELU',
32
+ 'Hardshrink', 'LeakyReLU', 'LogSigmoid', 'Softplus', 'Softshrink', 'PReLU',
33
+ 'Softsign', 'Tanhshrink', 'Softmin', 'Softmax', 'Softmax2d', 'LogSoftmax',
34
+
35
+ # others
36
+ 'Identity', 'SpikeBitwise',
37
+ ]
38
+
39
+
40
+ class Threshold(ElementWiseBlock):
41
+ r"""Thresholds each element of the input Tensor.
42
+
43
+ Threshold is defined as:
44
+
45
+ .. math::
46
+ y =
47
+ \begin{cases}
48
+ x, &\text{ if } x > \text{threshold} \\
49
+ \text{value}, &\text{ otherwise }
50
+ \end{cases}
51
+
52
+ Parameters
53
+ ----------
54
+ threshold : float
55
+ The value to threshold at.
56
+ value : float
57
+ The value to replace with.
58
+
59
+ Shape
60
+ -----
61
+ - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
62
+ - Output: :math:`(*)`, same shape as the input.
63
+
64
+ Examples
65
+ --------
66
+ .. code-block:: python
67
+
68
+ >>> import brainstate.nn as nn
69
+ >>> import brainstate
70
+ >>> m = nn.Threshold(0.1, 20)
71
+ >>> x = brainstate.random.randn(2)
72
+ >>> output = m(x)
73
+ """
74
+ __module__ = 'brainstate.nn'
75
+ threshold: float
76
+ value: float
77
+
78
+ def __init__(self, threshold: float, value: float) -> None:
79
+ super().__init__()
80
+ self.threshold = threshold
81
+ self.value = value
82
+
83
+ def __call__(self, x: ArrayLike) -> ArrayLike:
84
+ dtype = u.math.get_dtype(x)
85
+ return jnp.where(x > jnp.asarray(self.threshold, dtype=dtype),
86
+ x,
87
+ jnp.asarray(self.value, dtype=dtype))
88
+
89
+ def __repr__(self):
90
+ return f'{self.__class__.__name__}(threshold={self.threshold}, value={self.value})'
91
+
92
+
93
+ class ReLU(ElementWiseBlock):
94
+ r"""Applies the rectified linear unit function element-wise.
95
+
96
+ The ReLU function is defined as:
97
+
98
+ .. math::
99
+ \text{ReLU}(x) = (x)^+ = \max(0, x)
100
+
101
+ Shape
102
+ -----
103
+ - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
104
+ - Output: :math:`(*)`, same shape as the input.
105
+
106
+ Examples
107
+ --------
108
+ .. code-block:: python
109
+
110
+ >>> import brainstate.nn as nn
111
+ >>> import brainstate
112
+ >>> m = nn.ReLU()
113
+ >>> x = brainstate.random.randn(2)
114
+ >>> output = m(x)
115
+
116
+ An implementation of CReLU - https://arxiv.org/abs/1603.05201
117
+
118
+ .. code-block:: python
119
+
120
+ >>> import brainstate.nn as nn
121
+ >>> import brainstate
122
+ >>> import jax.numpy as jnp
123
+ >>> m = nn.ReLU()
124
+ >>> x = brainstate.random.randn(2).unsqueeze(0)
125
+ >>> output = jnp.concat((m(x), m(-x)))
126
+ """
127
+ __module__ = 'brainstate.nn'
128
+
129
+ def __call__(self, x: ArrayLike) -> ArrayLike:
130
+ return F.relu(x)
131
+
132
+ def __repr__(self):
133
+ return f'{self.__class__.__name__}()'
134
+
135
+
136
+ class RReLU(ElementWiseBlock):
137
+ r"""Applies the randomized leaky rectified liner unit function, element-wise.
138
+
139
+ As described in the paper `Empirical Evaluation of Rectified Activations in
140
+ Convolutional Network`_.
141
+
142
+ The function is defined as:
143
+
144
+ .. math::
145
+ \text{RReLU}(x) =
146
+ \begin{cases}
147
+ x & \text{if } x \geq 0 \\
148
+ ax & \text{ otherwise }
149
+ \end{cases}
150
+
151
+ where :math:`a` is randomly sampled from uniform distribution
152
+ :math:`\mathcal{U}(\text{lower}, \text{upper})`.
153
+
154
+ Parameters
155
+ ----------
156
+ lower : float, optional
157
+ Lower bound of the uniform distribution. Default: :math:`\frac{1}{8}`
158
+ upper : float, optional
159
+ Upper bound of the uniform distribution. Default: :math:`\frac{1}{3}`
160
+
161
+ Shape
162
+ -----
163
+ - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
164
+ - Output: :math:`(*)`, same shape as the input.
165
+
166
+ References
167
+ ----------
168
+ .. _`Empirical Evaluation of Rectified Activations in Convolutional Network`:
169
+ https://arxiv.org/abs/1505.00853
170
+
171
+ Examples
172
+ --------
173
+ .. code-block:: python
174
+
175
+ >>> import brainstate.nn as nn
176
+ >>> import brainstate
177
+ >>> m = nn.RReLU(0.1, 0.3)
178
+ >>> x = brainstate.random.randn(2)
179
+ >>> output = m(x)
180
+ """
181
+ __module__ = 'brainstate.nn'
182
+ lower: float
183
+ upper: float
184
+
185
+ def __init__(
186
+ self,
187
+ lower: float = 1. / 8,
188
+ upper: float = 1. / 3,
189
+ ):
190
+ super().__init__()
191
+ self.lower = lower
192
+ self.upper = upper
193
+
194
+ def __call__(self, x: ArrayLike) -> ArrayLike:
195
+ return F.rrelu(x, self.lower, self.upper)
196
+
197
+ def extra_repr(self):
198
+ return f'{self.__class__.__name__}(lower={self.lower}, upper={self.upper})'
199
+
200
+
201
+ class Hardtanh(ElementWiseBlock):
202
+ r"""Applies the HardTanh function element-wise.
203
+
204
+ HardTanh is defined as:
205
+
206
+ .. math::
207
+ \text{HardTanh}(x) = \begin{cases}
208
+ \text{max\_val} & \text{ if } x > \text{ max\_val } \\
209
+ \text{min\_val} & \text{ if } x < \text{ min\_val } \\
210
+ x & \text{ otherwise } \\
211
+ \end{cases}
212
+
213
+ Parameters
214
+ ----------
215
+ min_val : float, optional
216
+ Minimum value of the linear region range. Default: -1
217
+ max_val : float, optional
218
+ Maximum value of the linear region range. Default: 1
219
+
220
+ Notes
221
+ -----
222
+ Keyword arguments :attr:`min_value` and :attr:`max_value`
223
+ have been deprecated in favor of :attr:`min_val` and :attr:`max_val`.
224
+
225
+ Shape
226
+ -----
227
+ - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
228
+ - Output: :math:`(*)`, same shape as the input.
229
+
230
+ Examples
231
+ --------
232
+ .. code-block:: python
233
+
234
+ >>> import brainstate.nn as nn
235
+ >>> import brainstate
236
+ >>> m = nn.Hardtanh(-2, 2)
237
+ >>> x = brainstate.random.randn(2)
238
+ >>> output = m(x)
239
+ """
240
+ __module__ = 'brainstate.nn'
241
+ min_val: float
242
+ max_val: float
243
+
244
+ def __init__(
245
+ self,
246
+ min_val: float = -1.,
247
+ max_val: float = 1.,
248
+ ) -> None:
249
+ super().__init__()
250
+ self.min_val = min_val
251
+ self.max_val = max_val
252
+ assert self.max_val > self.min_val
253
+
254
+ def __call__(self, x: ArrayLike) -> ArrayLike:
255
+ return F.hard_tanh(x, self.min_val, self.max_val)
256
+
257
+ def extra_repr(self) -> str:
258
+ return f'{self.__class__.__name__}(min_val={self.min_val}, max_val={self.max_val})'
259
+
260
+
261
+ class ReLU6(Hardtanh):
262
+ r"""Applies the element-wise function.
263
+
264
+ ReLU6 is defined as:
265
+
266
+ .. math::
267
+ \text{ReLU6}(x) = \min(\max(0,x), 6)
268
+
269
+ Shape
270
+ -----
271
+ - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
272
+ - Output: :math:`(*)`, same shape as the input.
273
+
274
+ Examples
275
+ --------
276
+ .. code-block:: python
277
+
278
+ >>> import brainstate.nn as nn
279
+ >>> import brainstate
280
+ >>> m = nn.ReLU6()
281
+ >>> x = brainstate.random.randn(2)
282
+ >>> output = m(x)
283
+ """
284
+ __module__ = 'brainstate.nn'
285
+
286
+ def __init__(self):
287
+ super().__init__(0., 6.)
288
+
289
+
290
+ class Sigmoid(ElementWiseBlock):
291
+ r"""Applies the element-wise function.
292
+
293
+ Sigmoid is defined as:
294
+
295
+ .. math::
296
+ \text{Sigmoid}(x) = \sigma(x) = \frac{1}{1 + \exp(-x)}
297
+
298
+ Shape
299
+ -----
300
+ - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
301
+ - Output: :math:`(*)`, same shape as the input.
302
+
303
+ Examples
304
+ --------
305
+ .. code-block:: python
306
+
307
+ >>> import brainstate.nn as nn
308
+ >>> import brainstate
309
+ >>> m = nn.Sigmoid()
310
+ >>> x = brainstate.random.randn(2)
311
+ >>> output = m(x)
312
+ """
313
+ __module__ = 'brainstate.nn'
314
+
315
+ def __call__(self, x: ArrayLike) -> ArrayLike:
316
+ return F.sigmoid(x)
317
+
318
+
319
+ class Hardsigmoid(ElementWiseBlock):
320
+ r"""Applies the Hardsigmoid function element-wise.
321
+
322
+ Hardsigmoid is defined as:
323
+
324
+ .. math::
325
+ \text{Hardsigmoid}(x) = \begin{cases}
326
+ 0 & \text{if~} x \le -3, \\
327
+ 1 & \text{if~} x \ge +3, \\
328
+ x / 6 + 1 / 2 & \text{otherwise}
329
+ \end{cases}
330
+
331
+ Shape
332
+ -----
333
+ - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
334
+ - Output: :math:`(*)`, same shape as the input.
335
+
336
+ Examples
337
+ --------
338
+ .. code-block:: python
339
+
340
+ >>> import brainstate.nn as nn
341
+ >>> import brainstate
342
+ >>> m = nn.Hardsigmoid()
343
+ >>> x = brainstate.random.randn(2)
344
+ >>> output = m(x)
345
+ """
346
+ __module__ = 'brainstate.nn'
347
+
348
+ def __call__(self, x: ArrayLike) -> ArrayLike:
349
+ return F.hard_sigmoid(x)
350
+
351
+
352
+ class Tanh(ElementWiseBlock):
353
+ r"""Applies the Hyperbolic Tangent (Tanh) function element-wise.
354
+
355
+ Tanh is defined as:
356
+
357
+ .. math::
358
+ \text{Tanh}(x) = \tanh(x) = \frac{\exp(x) - \exp(-x)} {\exp(x) + \exp(-x)}
359
+
360
+ Shape
361
+ -----
362
+ - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
363
+ - Output: :math:`(*)`, same shape as the input.
364
+
365
+ Examples
366
+ --------
367
+ .. code-block:: python
368
+
369
+ >>> import brainstate.nn as nn
370
+ >>> import brainstate
371
+ >>> m = nn.Tanh()
372
+ >>> x = brainstate.random.randn(2)
373
+ >>> output = m(x)
374
+ """
375
+ __module__ = 'brainstate.nn'
376
+
377
+ def __call__(self, x: ArrayLike) -> ArrayLike:
378
+ return F.tanh(x)
379
+
380
+
381
+ class SiLU(ElementWiseBlock):
382
+ r"""Applies the Sigmoid Linear Unit (SiLU) function, element-wise.
383
+
384
+ The SiLU function is also known as the swish function.
385
+
386
+ .. math::
387
+ \text{silu}(x) = x * \sigma(x), \text{where } \sigma(x) \text{ is the logistic sigmoid.}
388
+
389
+ Notes
390
+ -----
391
+ See `Gaussian Error Linear Units (GELUs) <https://arxiv.org/abs/1606.08415>`_
392
+ where the SiLU (Sigmoid Linear Unit) was originally coined, and see
393
+ `Sigmoid-Weighted Linear Units for Neural Network Function Approximation
394
+ in Reinforcement Learning <https://arxiv.org/abs/1702.03118>`_ and `Swish:
395
+ a Self-Gated Activation Function <https://arxiv.org/abs/1710.05941v1>`_
396
+ where the SiLU was experimented with later.
397
+
398
+ Shape
399
+ -----
400
+ - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
401
+ - Output: :math:`(*)`, same shape as the input.
402
+
403
+ Examples
404
+ --------
405
+ .. code-block:: python
406
+
407
+ >>> import brainstate.nn as nn
408
+ >>> import brainstate
409
+ >>> m = nn.SiLU()
410
+ >>> x = brainstate.random.randn(2)
411
+ >>> output = m(x)
412
+ """
413
+ __module__ = 'brainstate.nn'
414
+
415
+ def __call__(self, x: ArrayLike) -> ArrayLike:
416
+ return F.silu(x)
417
+
418
+
419
+ class Mish(ElementWiseBlock):
420
+ r"""Applies the Mish function, element-wise.
421
+
422
+ Mish: A Self Regularized Non-Monotonic Neural Activation Function.
423
+
424
+ .. math::
425
+ \text{Mish}(x) = x * \text{Tanh}(\text{Softplus}(x))
426
+
427
+ Notes
428
+ -----
429
+ See `Mish: A Self Regularized Non-Monotonic Neural Activation Function
430
+ <https://arxiv.org/abs/1908.08681>`_
431
+
432
+ Shape
433
+ -----
434
+ - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
435
+ - Output: :math:`(*)`, same shape as the input.
436
+
437
+ Examples
438
+ --------
439
+ .. code-block:: python
440
+
441
+ >>> import brainstate.nn as nn
442
+ >>> import brainstate
443
+ >>> m = nn.Mish()
444
+ >>> x = brainstate.random.randn(2)
445
+ >>> output = m(x)
446
+ """
447
+ __module__ = 'brainstate.nn'
448
+
449
+ def __call__(self, x: ArrayLike) -> ArrayLike:
450
+ return F.mish(x)
451
+
452
+
453
+ class Hardswish(ElementWiseBlock):
454
+ r"""Applies the Hardswish function, element-wise.
455
+
456
+ As described in the paper `Searching for MobileNetV3
457
+ <https://arxiv.org/abs/1905.02244>`_.
458
+
459
+ Hardswish is defined as:
460
+
461
+ .. math::
462
+ \text{Hardswish}(x) = \begin{cases}
463
+ 0 & \text{if~} x \le -3, \\
464
+ x & \text{if~} x \ge +3, \\
465
+ x \cdot (x + 3) /6 & \text{otherwise}
466
+ \end{cases}
467
+
468
+ Shape
469
+ -----
470
+ - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
471
+ - Output: :math:`(*)`, same shape as the input.
472
+
473
+ Examples
474
+ --------
475
+ .. code-block:: python
476
+
477
+ >>> import brainstate.nn as nn
478
+ >>> import brainstate
479
+ >>> m = nn.Hardswish()
480
+ >>> x = brainstate.random.randn(2)
481
+ >>> output = m(x)
482
+ """
483
+ __module__ = 'brainstate.nn'
484
+
485
+ def __call__(self, x: ArrayLike) -> ArrayLike:
486
+ return F.hard_swish(x)
487
+
488
+
489
+ class ELU(ElementWiseBlock):
490
+ r"""Applies the Exponential Linear Unit (ELU) function, element-wise.
491
+
492
+ As described in the paper: `Fast and Accurate Deep Network Learning by
493
+ Exponential Linear Units (ELUs) <https://arxiv.org/abs/1511.07289>`__.
494
+
495
+ ELU is defined as:
496
+
497
+ .. math::
498
+ \text{ELU}(x) = \begin{cases}
499
+ x, & \text{ if } x > 0\\
500
+ \alpha * (\exp(x) - 1), & \text{ if } x \leq 0
501
+ \end{cases}
502
+
503
+ Parameters
504
+ ----------
505
+ alpha : float, optional
506
+ The :math:`\alpha` value for the ELU formulation. Default: 1.0
507
+
508
+ Shape
509
+ -----
510
+ - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
511
+ - Output: :math:`(*)`, same shape as the input.
512
+
513
+ Examples
514
+ --------
515
+ .. code-block:: python
516
+
517
+ >>> import brainstate.nn as nn
518
+ >>> import brainstate
519
+ >>> m = nn.ELU()
520
+ >>> x = brainstate.random.randn(2)
521
+ >>> output = m(x)
522
+ """
523
+ __module__ = 'brainstate.nn'
524
+ alpha: float
525
+
526
+ def __init__(self, alpha: float = 1.) -> None:
527
+ super().__init__()
528
+ self.alpha = alpha
529
+
530
+ def __call__(self, x: ArrayLike) -> ArrayLike:
531
+ return F.elu(x, self.alpha)
532
+
533
+ def extra_repr(self) -> str:
534
+ return f'{self.__class__.__name__}(alpha={self.alpha})'
535
+
536
+
537
+ class CELU(ElementWiseBlock):
538
+ r"""Applies the element-wise function.
539
+
540
+ .. math::
541
+ \text{CELU}(x) = \max(0,x) + \min(0, \alpha * (\exp(x/\alpha) - 1))
542
+
543
+ More details can be found in the paper `Continuously Differentiable Exponential
544
+ Linear Units`_ .
545
+
546
+ Parameters
547
+ ----------
548
+ alpha : float, optional
549
+ The :math:`\alpha` value for the CELU formulation. Default: 1.0
550
+
551
+ Shape
552
+ -----
553
+ - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
554
+ - Output: :math:`(*)`, same shape as the input.
555
+
556
+ References
557
+ ----------
558
+ .. _`Continuously Differentiable Exponential Linear Units`:
559
+ https://arxiv.org/abs/1704.07483
560
+
561
+ Examples
562
+ --------
563
+ .. code-block:: python
564
+
565
+ >>> import brainstate.nn as nn
566
+ >>> import brainstate
567
+ >>> m = nn.CELU()
568
+ >>> x = brainstate.random.randn(2)
569
+ >>> output = m(x)
570
+ """
571
+ __module__ = 'brainstate.nn'
572
+ alpha: float
573
+
574
+ def __init__(self, alpha: float = 1.) -> None:
575
+ super().__init__()
576
+ self.alpha = alpha
577
+
578
+ def __call__(self, x: ArrayLike) -> ArrayLike:
579
+ return F.celu(x, self.alpha)
580
+
581
+ def extra_repr(self) -> str:
582
+ return f'{self.__class__.__name__}(alpha={self.alpha})'
583
+
584
+
585
+ class SELU(ElementWiseBlock):
586
+ r"""Applied element-wise.
587
+
588
+ .. math::
589
+ \text{SELU}(x) = \text{scale} * (\max(0,x) + \min(0, \alpha * (\exp(x) - 1)))
590
+
591
+ with :math:`\alpha = 1.6732632423543772848170429916717` and
592
+ :math:`\text{scale} = 1.0507009873554804934193349852946`.
593
+
594
+ More details can be found in the paper `Self-Normalizing Neural Networks`_ .
595
+
596
+ Shape
597
+ -----
598
+ - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
599
+ - Output: :math:`(*)`, same shape as the input.
600
+
601
+ References
602
+ ----------
603
+ .. _Self-Normalizing Neural Networks: https://arxiv.org/abs/1706.02515
604
+
605
+ Examples
606
+ --------
607
+ .. code-block:: python
608
+
609
+ >>> import brainstate.nn as nn
610
+ >>> import brainstate
611
+ >>> m = nn.SELU()
612
+ >>> x = brainstate.random.randn(2)
613
+ >>> output = m(x)
614
+ """
615
+ __module__ = 'brainstate.nn'
616
+
617
+ def __call__(self, x: ArrayLike) -> ArrayLike:
618
+ return F.selu(x)
619
+
620
+
621
+ class GLU(ElementWiseBlock):
622
+ r"""Applies the gated linear unit function.
623
+
624
+ .. math::
625
+ {GLU}(a, b)= a \otimes \sigma(b)
626
+
627
+ where :math:`a` is the first half of the input matrices and :math:`b` is
628
+ the second half.
629
+
630
+ Parameters
631
+ ----------
632
+ dim : int, optional
633
+ The dimension on which to split the input. Default: -1
634
+
635
+ Shape
636
+ -----
637
+ - Input: :math:`(\ast_1, N, \ast_2)` where `*` means, any number of additional
638
+ dimensions
639
+ - Output: :math:`(\ast_1, M, \ast_2)` where :math:`M=N/2`
640
+
641
+ Examples
642
+ --------
643
+ .. code-block:: python
644
+
645
+ >>> import brainstate.nn as nn
646
+ >>> import brainstate
647
+ >>> m = nn.GLU()
648
+ >>> x = brainstate.random.randn(4, 2)
649
+ >>> output = m(x)
650
+ """
651
+ __module__ = 'brainstate.nn'
652
+ dim: int
653
+
654
+ def __init__(self, dim: int = -1) -> None:
655
+ super().__init__()
656
+ self.dim = dim
657
+
658
+ def __call__(self, x: ArrayLike) -> ArrayLike:
659
+ return F.glu(x, self.dim)
660
+
661
+ def __repr__(self):
662
+ return f'{self.__class__.__name__}(dim={self.dim})'
663
+
664
+
665
+ class GELU(ElementWiseBlock):
666
+ r"""Applies the Gaussian Error Linear Units function.
667
+
668
+ .. math:: \text{GELU}(x) = x * \Phi(x)
669
+
670
+ where :math:`\Phi(x)` is the Cumulative Distribution Function for Gaussian
671
+ Distribution.
672
+
673
+ When the approximate argument is True, Gelu is estimated with:
674
+
675
+ .. math:: \text{GELU}(x) = 0.5 * x * (1 + \text{Tanh}(\sqrt(2 / \pi) * (x + 0.044715 * x^3)))
676
+
677
+ Parameters
678
+ ----------
679
+ approximate : bool, optional
680
+ Whether to use the tanh approximation algorithm. Default: False
681
+
682
+ Shape
683
+ -----
684
+ - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
685
+ - Output: :math:`(*)`, same shape as the input.
686
+
687
+ Examples
688
+ --------
689
+ .. code-block:: python
690
+
691
+ >>> import brainstate.nn as nn
692
+ >>> import brainstate
693
+ >>> m = nn.GELU()
694
+ >>> x = brainstate.random.randn(2)
695
+ >>> output = m(x)
696
+ """
697
+ __module__ = 'brainstate.nn'
698
+ approximate: bool
699
+
700
+ def __init__(self, approximate: bool = False) -> None:
701
+ super().__init__()
702
+ self.approximate = approximate
703
+
704
+ def __call__(self, x: ArrayLike) -> ArrayLike:
705
+ return F.gelu(x, approximate=self.approximate)
706
+
707
+ def __repr__(self):
708
+ return f'{self.__class__.__name__}(approximate={self.approximate})'
709
+
710
+
711
+ class Hardshrink(ElementWiseBlock):
712
+ r"""Applies the Hard Shrinkage (Hardshrink) function element-wise.
713
+
714
+ Hardshrink is defined as:
715
+
716
+ .. math::
717
+ \text{HardShrink}(x) =
718
+ \begin{cases}
719
+ x, & \text{ if } x > \lambda \\
720
+ x, & \text{ if } x < -\lambda \\
721
+ 0, & \text{ otherwise }
722
+ \end{cases}
723
+
724
+ Parameters
725
+ ----------
726
+ lambd : float, optional
727
+ The :math:`\lambda` value for the Hardshrink formulation. Default: 0.5
728
+
729
+ Shape
730
+ -----
731
+ - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
732
+ - Output: :math:`(*)`, same shape as the input.
733
+
734
+ Examples
735
+ --------
736
+ .. code-block:: python
737
+
738
+ >>> import brainstate.nn as nn
739
+ >>> import brainstate
740
+ >>> m = nn.Hardshrink()
741
+ >>> x = brainstate.random.randn(2)
742
+ >>> output = m(x)
743
+ """
744
+ __module__ = 'brainstate.nn'
745
+ lambd: float
746
+
747
+ def __init__(self, lambd: float = 0.5) -> None:
748
+ super().__init__()
749
+ self.lambd = lambd
750
+
751
+ def __call__(self, x: ArrayLike) -> ArrayLike:
752
+ return F.hard_shrink(x, self.lambd)
753
+
754
+ def __repr__(self):
755
+ return f'{self.__class__.__name__}(lambd={self.lambd})'
756
+
757
+
758
+ class LeakyReLU(ElementWiseBlock):
759
+ r"""Applies the element-wise function.
760
+
761
+ .. math::
762
+ \text{LeakyReLU}(x) = \max(0, x) + \text{negative\_slope} * \min(0, x)
763
+
764
+ or
765
+
766
+ .. math::
767
+ \text{LeakyReLU}(x) =
768
+ \begin{cases}
769
+ x, & \text{ if } x \geq 0 \\
770
+ \text{negative\_slope} \times x, & \text{ otherwise }
771
+ \end{cases}
772
+
773
+ Parameters
774
+ ----------
775
+ negative_slope : float, optional
776
+ Controls the angle of the negative slope (which is used for
777
+ negative input values). Default: 1e-2
778
+
779
+ Shape
780
+ -----
781
+ - Input: :math:`(*)` where `*` means, any number of additional
782
+ dimensions
783
+ - Output: :math:`(*)`, same shape as the input
784
+
785
+ Examples
786
+ --------
787
+ .. code-block:: python
788
+
789
+ >>> import brainstate.nn as nn
790
+ >>> import brainstate
791
+ >>> m = nn.LeakyReLU(0.1)
792
+ >>> x = brainstate.random.randn(2)
793
+ >>> output = m(x)
794
+ """
795
+ __module__ = 'brainstate.nn'
796
+ negative_slope: float
797
+
798
+ def __init__(self, negative_slope: float = 1e-2) -> None:
799
+ super().__init__()
800
+ self.negative_slope = negative_slope
801
+
802
+ def __call__(self, x: ArrayLike) -> ArrayLike:
803
+ return F.leaky_relu(x, self.negative_slope)
804
+
805
+ def __repr__(self):
806
+ return f'{self.__class__.__name__}(negative_slope={self.negative_slope})'
807
+
808
+
809
+ class LogSigmoid(ElementWiseBlock):
810
+ r"""Applies the element-wise function.
811
+
812
+ .. math::
813
+ \text{LogSigmoid}(x) = \log\left(\frac{ 1 }{ 1 + \exp(-x)}\right)
814
+
815
+ Shape
816
+ -----
817
+ - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
818
+ - Output: :math:`(*)`, same shape as the input.
819
+
820
+ Examples
821
+ --------
822
+ .. code-block:: python
823
+
824
+ >>> import brainstate.nn as nn
825
+ >>> import brainstate
826
+ >>> m = nn.LogSigmoid()
827
+ >>> x = brainstate.random.randn(2)
828
+ >>> output = m(x)
829
+ """
830
+ __module__ = 'brainstate.nn'
831
+
832
+ def __call__(self, x: ArrayLike) -> ArrayLike:
833
+ return F.log_sigmoid(x)
834
+
835
+
836
+ class Softplus(ElementWiseBlock):
837
+ r"""Applies the Softplus function element-wise.
838
+
839
+ .. math::
840
+ \text{Softplus}(x) = \frac{1}{\beta} * \log(1 + \exp(\beta * x))
841
+
842
+ SoftPlus is a smooth approximation to the ReLU function and can be used
843
+ to constrain the output of a machine to always be positive.
844
+
845
+ For numerical stability the implementation reverts to the linear function
846
+ when :math:`input \times \beta > threshold`.
847
+
848
+ Shape
849
+ -----
850
+ - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
851
+ - Output: :math:`(*)`, same shape as the input.
852
+
853
+ Examples
854
+ --------
855
+ .. code-block:: python
856
+
857
+ >>> import brainstate.nn as nn
858
+ >>> import brainstate
859
+ >>> m = nn.Softplus()
860
+ >>> x = brainstate.random.randn(2)
861
+ >>> output = m(x)
862
+ """
863
+ __module__ = 'brainstate.nn'
864
+
865
+ def __call__(self, x: ArrayLike) -> ArrayLike:
866
+ return F.softplus(x)
867
+
868
+
869
+ class Softshrink(ElementWiseBlock):
870
+ r"""Applies the soft shrinkage function elementwise.
871
+
872
+ .. math::
873
+ \text{SoftShrinkage}(x) =
874
+ \begin{cases}
875
+ x - \lambda, & \text{ if } x > \lambda \\
876
+ x + \lambda, & \text{ if } x < -\lambda \\
877
+ 0, & \text{ otherwise }
878
+ \end{cases}
879
+
880
+ Parameters
881
+ ----------
882
+ lambd : float, optional
883
+ The :math:`\lambda` (must be no less than zero) value for the
884
+ Softshrink formulation. Default: 0.5
885
+
886
+ Shape
887
+ -----
888
+ - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
889
+ - Output: :math:`(*)`, same shape as the input.
890
+
891
+ Examples
892
+ --------
893
+ .. code-block:: python
894
+
895
+ >>> import brainstate.nn as nn
896
+ >>> import brainstate
897
+ >>> m = nn.Softshrink()
898
+ >>> x = brainstate.random.randn(2)
899
+ >>> output = m(x)
900
+ """
901
+ __module__ = 'brainstate.nn'
902
+ lambd: float
903
+
904
+ def __init__(self, lambd: float = 0.5) -> None:
905
+ super().__init__()
906
+ self.lambd = lambd
907
+
908
+ def __call__(self, x: ArrayLike) -> ArrayLike:
909
+ return F.soft_shrink(x, self.lambd)
910
+
911
+ def __repr__(self):
912
+ return f'{self.__class__.__name__}(lambd={self.lambd})'
913
+
914
+
915
+ class PReLU(ElementWiseBlock):
916
+ r"""Applies the element-wise function.
917
+
918
+ .. math::
919
+ \text{PReLU}(x) = \max(0,x) + a * \min(0,x)
920
+
921
+ or
922
+
923
+ .. math::
924
+ \text{PReLU}(x) =
925
+ \begin{cases}
926
+ x, & \text{ if } x \geq 0 \\
927
+ ax, & \text{ otherwise }
928
+ \end{cases}
929
+
930
+ Here :math:`a` is a learnable parameter. When called without arguments,
931
+ `nn.PReLU()` uses a single parameter :math:`a` across all input channels.
932
+ If called with `nn.PReLU(nChannels)`, a separate :math:`a` is used for
933
+ each input channel.
934
+
935
+ Parameters
936
+ ----------
937
+ num_parameters : int, optional
938
+ Number of :math:`a` to learn. Although it takes an int as input,
939
+ there is only two values are legitimate: 1, or the number of channels
940
+ at input. Default: 1
941
+ init : float, optional
942
+ The initial value of :math:`a`. Default: 0.25
943
+ dtype : optional
944
+ The data type for the weight parameter.
945
+
946
+ Shape
947
+ -----
948
+ - Input: :math:`( *)` where `*` means, any number of additional dimensions.
949
+ - Output: :math:`(*)`, same shape as the input.
950
+
951
+ Attributes
952
+ ----------
953
+ weight : Tensor
954
+ The learnable weights of shape (:attr:`num_parameters`).
955
+
956
+ Notes
957
+ -----
958
+ - Weight decay should not be used when learning :math:`a` for good performance.
959
+ - Channel dim is the 2nd dim of input. When input has dims < 2, then there is
960
+ no channel dim and the number of channels = 1.
961
+
962
+ Examples
963
+ --------
964
+ .. code-block:: python
965
+
966
+ >>> import brainstate
967
+ >>> m = brainstate.nn.PReLU()
968
+ >>> x = brainstate.random.randn(2)
969
+ >>> output = m(x)
970
+ """
971
+ __module__ = 'brainstate.nn'
972
+ num_parameters: int
973
+
974
+ def __init__(self, num_parameters: int = 1, init: float = 0.25, dtype=None) -> None:
975
+ super().__init__()
976
+ self.num_parameters = num_parameters
977
+ self.weight = ParamState(jnp.ones(num_parameters, dtype=dtype) * init)
978
+
979
+ def __call__(self, x: ArrayLike) -> ArrayLike:
980
+ return F.prelu(x, self.weight.value)
981
+
982
+ def __repr__(self):
983
+ return f'{self.__class__.__name__}(num_parameters={self.num_parameters})'
984
+
985
+
986
+ class Softsign(ElementWiseBlock):
987
+ r"""Applies the element-wise function.
988
+
989
+ .. math::
990
+ \text{SoftSign}(x) = \frac{x}{ 1 + |x|}
991
+
992
+ Shape
993
+ -----
994
+ - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
995
+ - Output: :math:`(*)`, same shape as the input.
996
+
997
+ Examples
998
+ --------
999
+ .. code-block:: python
1000
+
1001
+ >>> import brainstate.nn as nn
1002
+ >>> import brainstate
1003
+ >>> m = nn.Softsign()
1004
+ >>> x = brainstate.random.randn(2)
1005
+ >>> output = m(x)
1006
+ """
1007
+ __module__ = 'brainstate.nn'
1008
+
1009
+ def __call__(self, x: ArrayLike) -> ArrayLike:
1010
+ return F.soft_sign(x)
1011
+
1012
+
1013
+ class Tanhshrink(ElementWiseBlock):
1014
+ r"""Applies the element-wise function.
1015
+
1016
+ .. math::
1017
+ \text{Tanhshrink}(x) = x - \tanh(x)
1018
+
1019
+ Shape
1020
+ -----
1021
+ - Input: :math:`(*)`, where :math:`*` means any number of dimensions.
1022
+ - Output: :math:`(*)`, same shape as the input.
1023
+
1024
+ Examples
1025
+ --------
1026
+ .. code-block:: python
1027
+
1028
+ >>> import brainstate.nn as nn
1029
+ >>> import brainstate
1030
+ >>> m = nn.Tanhshrink()
1031
+ >>> x = brainstate.random.randn(2)
1032
+ >>> output = m(x)
1033
+ """
1034
+ __module__ = 'brainstate.nn'
1035
+
1036
+ def __call__(self, x: ArrayLike) -> ArrayLike:
1037
+ return F.tanh_shrink(x)
1038
+
1039
+
1040
+ class Softmin(ElementWiseBlock):
1041
+ r"""Applies the Softmin function to an n-dimensional input Tensor.
1042
+
1043
+ Rescales the input so that the elements of the n-dimensional output Tensor
1044
+ lie in the range `[0, 1]` and sum to 1.
1045
+
1046
+ Softmin is defined as:
1047
+
1048
+ .. math::
1049
+ \text{Softmin}(x_{i}) = \frac{\exp(-x_i)}{\sum_j \exp(-x_j)}
1050
+
1051
+ Parameters
1052
+ ----------
1053
+ dim : int, optional
1054
+ A dimension along which Softmin will be computed (so every slice
1055
+ along dim will sum to 1).
1056
+
1057
+ Shape
1058
+ -----
1059
+ - Input: :math:`(*)` where `*` means, any number of additional dimensions
1060
+ - Output: :math:`(*)`, same shape as the input
1061
+
1062
+ Returns
1063
+ -------
1064
+ Tensor
1065
+ A Tensor of the same dimension and shape as the input, with
1066
+ values in the range [0, 1]
1067
+
1068
+ Examples
1069
+ --------
1070
+ .. code-block:: python
1071
+
1072
+ >>> import brainstate.nn as nn
1073
+ >>> import brainstate
1074
+ >>> m = nn.Softmin(dim=1)
1075
+ >>> x = brainstate.random.randn(2, 3)
1076
+ >>> output = m(x)
1077
+ """
1078
+ __module__ = 'brainstate.nn'
1079
+ dim: Optional[int]
1080
+
1081
+ def __init__(self, dim: Optional[int] = None) -> None:
1082
+ super().__init__()
1083
+ self.dim = dim
1084
+
1085
+ def __call__(self, x: ArrayLike) -> ArrayLike:
1086
+ return F.softmin(x, self.dim)
1087
+
1088
+ def __repr__(self):
1089
+ return f'{self.__class__.__name__}(dim={self.dim})'
1090
+
1091
+
1092
+ class Softmax(ElementWiseBlock):
1093
+ r"""Applies the Softmax function to an n-dimensional input Tensor.
1094
+
1095
+ Rescales the input so that the elements of the n-dimensional output Tensor
1096
+ lie in the range [0,1] and sum to 1.
1097
+
1098
+ Softmax is defined as:
1099
+
1100
+ .. math::
1101
+ \text{Softmax}(x_{i}) = \frac{\exp(x_i)}{\sum_j \exp(x_j)}
1102
+
1103
+ When the input Tensor is a sparse tensor then the unspecified
1104
+ values are treated as ``-inf``.
1105
+
1106
+ Parameters
1107
+ ----------
1108
+ dim : int, optional
1109
+ A dimension along which Softmax will be computed (so every slice
1110
+ along dim will sum to 1).
1111
+
1112
+ Shape
1113
+ -----
1114
+ - Input: :math:`(*)` where `*` means, any number of additional dimensions
1115
+ - Output: :math:`(*)`, same shape as the input
1116
+
1117
+ Returns
1118
+ -------
1119
+ Tensor
1120
+ A Tensor of the same dimension and shape as the input with
1121
+ values in the range [0, 1]
1122
+
1123
+ Notes
1124
+ -----
1125
+ This module doesn't work directly with NLLLoss, which expects the Log to be
1126
+ computed between the Softmax and itself. Use `LogSoftmax` instead (it's
1127
+ faster and has better numerical properties).
1128
+
1129
+ Examples
1130
+ --------
1131
+ .. code-block:: python
1132
+
1133
+ >>> import brainstate.nn as nn
1134
+ >>> import brainstate
1135
+ >>> m = nn.Softmax(dim=1)
1136
+ >>> x = brainstate.random.randn(2, 3)
1137
+ >>> output = m(x)
1138
+ """
1139
+ __module__ = 'brainstate.nn'
1140
+ dim: Optional[int]
1141
+
1142
+ def __init__(self, dim: Optional[int] = None) -> None:
1143
+ super().__init__()
1144
+ self.dim = dim
1145
+
1146
+ def __call__(self, x: ArrayLike) -> ArrayLike:
1147
+ return F.softmax(x, self.dim)
1148
+
1149
+ def __repr__(self):
1150
+ return f'{self.__class__.__name__}(dim={self.dim})'
1151
+
1152
+
1153
+ class Softmax2d(ElementWiseBlock):
1154
+ r"""Applies SoftMax over features to each spatial location.
1155
+
1156
+ When given an image of ``Channels x Height x Width``, it will
1157
+ apply `Softmax` to each location :math:`(Channels, h_i, w_j)`
1158
+
1159
+ Shape
1160
+ -----
1161
+ - Input: :math:`(N, C, H, W)` or :math:`(C, H, W)`.
1162
+ - Output: :math:`(N, C, H, W)` or :math:`(C, H, W)` (same shape as input)
1163
+
1164
+ Returns
1165
+ -------
1166
+ Tensor
1167
+ A Tensor of the same dimension and shape as the input with
1168
+ values in the range [0, 1]
1169
+
1170
+ Examples
1171
+ --------
1172
+ .. code-block:: python
1173
+
1174
+ >>> import brainstate.nn as nn
1175
+ >>> import brainstate
1176
+ >>> m = nn.Softmax2d()
1177
+ >>> # you softmax over the 2nd dimension
1178
+ >>> x = brainstate.random.randn(2, 3, 12, 13)
1179
+ >>> output = m(x)
1180
+ """
1181
+ __module__ = 'brainstate.nn'
1182
+
1183
+ def __call__(self, x: ArrayLike) -> ArrayLike:
1184
+ assert x.ndim == 4 or x.ndim == 3, 'Softmax2d requires a 3D or 4D tensor as input'
1185
+ return F.softmax(x, -3)
1186
+
1187
+
1188
+ class LogSoftmax(ElementWiseBlock):
1189
+ r"""Applies the :math:`\log(\text{Softmax}(x))` function to an n-dimensional input Tensor.
1190
+
1191
+ The LogSoftmax formulation can be simplified as:
1192
+
1193
+ .. math::
1194
+ \text{LogSoftmax}(x_{i}) = \log\left(\frac{\exp(x_i) }{ \sum_j \exp(x_j)} \right)
1195
+
1196
+ Parameters
1197
+ ----------
1198
+ dim : int, optional
1199
+ A dimension along which LogSoftmax will be computed.
1200
+
1201
+ Shape
1202
+ -----
1203
+ - Input: :math:`(*)` where `*` means, any number of additional dimensions
1204
+ - Output: :math:`(*)`, same shape as the input
1205
+
1206
+ Returns
1207
+ -------
1208
+ Tensor
1209
+ A Tensor of the same dimension and shape as the input with
1210
+ values in the range [-inf, 0)
1211
+
1212
+ Examples
1213
+ --------
1214
+ .. code-block:: python
1215
+
1216
+ >>> import brainstate.nn as nn
1217
+ >>> import brainstate
1218
+ >>> m = nn.LogSoftmax(dim=1)
1219
+ >>> x = brainstate.random.randn(2, 3)
1220
+ >>> output = m(x)
1221
+ """
1222
+ __module__ = 'brainstate.nn'
1223
+ dim: Optional[int]
1224
+
1225
+ def __init__(self, dim: Optional[int] = None) -> None:
1226
+ super().__init__()
1227
+ self.dim = dim
1228
+
1229
+ def __call__(self, x: ArrayLike) -> ArrayLike:
1230
+ return F.log_softmax(x, self.dim)
1231
+
1232
+ def __repr__(self):
1233
+ return f'{self.__class__.__name__}(dim={self.dim})'
1234
+
1235
+
1236
+ class Identity(ElementWiseBlock):
1237
+ r"""A placeholder identity operator that is argument-insensitive.
1238
+
1239
+ Examples
1240
+ --------
1241
+ .. code-block:: python
1242
+
1243
+ >>> import brainstate.nn as nn
1244
+ >>> m = nn.Identity()
1245
+ >>> x = brainstate.random.randn(2, 3)
1246
+ >>> output = m(x)
1247
+ >>> assert (output == x).all()
1248
+ """
1249
+ __module__ = 'brainstate.nn'
1250
+
1251
+ def __call__(self, x):
1252
+ return x
1253
+
1254
+
1255
+ class SpikeBitwise(ElementWiseBlock):
1256
+ r"""Bitwise addition for the spiking inputs.
1257
+
1258
+ .. math::
1259
+
1260
+ \begin{array}{ccc}
1261
+ \hline \text { Mode } & \text { Expression for } \mathrm{g}(\mathrm{x}, \mathrm{y}) & \text { Code for } \mathrm{g}(\mathrm{x}, \mathrm{y}) \\
1262
+ \hline \text { ADD } & x+y & x+y \\
1263
+ \text { AND } & x \cap y & x \cdot y \\
1264
+ \text { IAND } & (\neg x) \cap y & (1-x) \cdot y \\
1265
+ \text { OR } & x \cup y & (x+y)-(x \cdot y) \\
1266
+ \hline
1267
+ \end{array}
1268
+
1269
+ Parameters
1270
+ ----------
1271
+ op : str, optional
1272
+ The bitwise operation. Default: 'add'
1273
+ name : str, optional
1274
+ The name of the dynamic system.
1275
+
1276
+ Examples
1277
+ --------
1278
+ .. code-block:: python
1279
+
1280
+ >>> import brainstate.nn as nn
1281
+ >>> m = nn.SpikeBitwise(op='and')
1282
+ >>> x = brainstate.random.randn(2, 3) > 0
1283
+ >>> y = brainstate.random.randn(2, 3) > 0
1284
+ >>> output = m(x, y)
1285
+ """
1286
+ __module__ = 'brainstate.nn'
1287
+
1288
+ def __init__(
1289
+ self,
1290
+ op: str = 'add',
1291
+ name: Optional[str] = None
1292
+ ) -> None:
1293
+ super().__init__(name=name)
1294
+ self.op = op
1295
+
1296
+ def __call__(self, x, y):
1297
+ import braintools
1298
+ return braintools.spike_bitwise(x, y, self.op)