brainstate 0.1.9__py2.py3-none-any.whl → 0.2.0__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.
- brainstate/__init__.py +130 -19
- brainstate/_compatible_import.py +201 -9
- brainstate/_compatible_import_test.py +681 -0
- brainstate/_deprecation.py +210 -0
- brainstate/_deprecation_test.py +2319 -0
- brainstate/{util/error.py → _error.py} +10 -20
- brainstate/_state.py +94 -47
- brainstate/_state_test.py +1 -1
- brainstate/_utils.py +1 -1
- brainstate/environ.py +1279 -347
- brainstate/environ_test.py +1187 -26
- brainstate/graph/__init__.py +6 -13
- brainstate/graph/_node.py +240 -0
- brainstate/graph/_node_test.py +589 -0
- brainstate/graph/{_graph_operation.py → _operation.py} +632 -746
- brainstate/graph/_operation_test.py +1147 -0
- brainstate/mixin.py +1209 -141
- brainstate/mixin_test.py +991 -51
- brainstate/nn/__init__.py +74 -72
- brainstate/nn/_activations.py +587 -295
- brainstate/nn/_activations_test.py +109 -86
- brainstate/nn/_collective_ops.py +393 -274
- brainstate/nn/_collective_ops_test.py +746 -15
- brainstate/nn/_common.py +114 -66
- brainstate/nn/_common_test.py +154 -0
- brainstate/nn/_conv.py +1652 -143
- brainstate/nn/_conv_test.py +838 -227
- brainstate/nn/_delay.py +95 -29
- brainstate/nn/_delay_test.py +25 -20
- brainstate/nn/_dropout.py +359 -167
- brainstate/nn/_dropout_test.py +429 -52
- brainstate/nn/_dynamics.py +14 -90
- brainstate/nn/_dynamics_test.py +1 -12
- brainstate/nn/_elementwise.py +492 -313
- brainstate/nn/_elementwise_test.py +806 -145
- brainstate/nn/_embedding.py +369 -19
- brainstate/nn/_embedding_test.py +156 -0
- brainstate/nn/{_fixedprob.py → _event_fixedprob.py} +10 -16
- brainstate/nn/{_fixedprob_test.py → _event_fixedprob_test.py} +6 -5
- brainstate/nn/{_linear_mv.py → _event_linear.py} +2 -2
- brainstate/nn/{_linear_mv_test.py → _event_linear_test.py} +6 -5
- brainstate/nn/_exp_euler.py +200 -38
- brainstate/nn/_exp_euler_test.py +350 -8
- brainstate/nn/_linear.py +391 -71
- brainstate/nn/_linear_test.py +427 -59
- brainstate/nn/_metrics.py +1070 -0
- brainstate/nn/_metrics_test.py +611 -0
- brainstate/nn/_module.py +10 -3
- brainstate/nn/_module_test.py +1 -1
- brainstate/nn/_normalizations.py +688 -329
- brainstate/nn/_normalizations_test.py +663 -37
- brainstate/nn/_paddings.py +1020 -0
- brainstate/nn/_paddings_test.py +723 -0
- brainstate/nn/_poolings.py +1404 -342
- brainstate/nn/_poolings_test.py +828 -92
- brainstate/nn/{_rate_rnns.py → _rnns.py} +446 -54
- brainstate/nn/_rnns_test.py +593 -0
- brainstate/nn/_utils.py +132 -5
- brainstate/nn/_utils_test.py +402 -0
- brainstate/{init/_random_inits.py → nn/init.py} +301 -45
- brainstate/{init/_random_inits_test.py → nn/init_test.py} +51 -20
- brainstate/random/__init__.py +247 -1
- brainstate/random/_rand_funs.py +668 -346
- brainstate/random/_rand_funs_test.py +74 -1
- brainstate/random/_rand_seed.py +541 -76
- brainstate/random/_rand_seed_test.py +1 -1
- brainstate/random/_rand_state.py +601 -393
- brainstate/random/_rand_state_test.py +551 -0
- brainstate/transform/__init__.py +59 -0
- brainstate/transform/_ad_checkpoint.py +176 -0
- brainstate/{compile → transform}/_ad_checkpoint_test.py +1 -1
- brainstate/{augment → transform}/_autograd.py +360 -113
- brainstate/{augment → transform}/_autograd_test.py +2 -2
- brainstate/transform/_conditions.py +316 -0
- brainstate/{compile → transform}/_conditions_test.py +11 -11
- brainstate/{compile → transform}/_error_if.py +22 -20
- brainstate/{compile → transform}/_error_if_test.py +1 -1
- brainstate/transform/_eval_shape.py +145 -0
- brainstate/{augment → transform}/_eval_shape_test.py +1 -1
- brainstate/{compile → transform}/_jit.py +99 -46
- brainstate/{compile → transform}/_jit_test.py +3 -3
- brainstate/{compile → transform}/_loop_collect_return.py +219 -80
- brainstate/{compile → transform}/_loop_collect_return_test.py +1 -1
- brainstate/{compile → transform}/_loop_no_collection.py +133 -34
- brainstate/{compile → transform}/_loop_no_collection_test.py +2 -2
- brainstate/transform/_make_jaxpr.py +2016 -0
- brainstate/transform/_make_jaxpr_test.py +1510 -0
- brainstate/transform/_mapping.py +529 -0
- brainstate/transform/_mapping_test.py +194 -0
- brainstate/{compile → transform}/_progress_bar.py +78 -25
- brainstate/{augment → transform}/_random.py +65 -45
- brainstate/{compile → transform}/_unvmap.py +102 -5
- brainstate/transform/_util.py +286 -0
- brainstate/typing.py +594 -61
- brainstate/typing_test.py +780 -0
- brainstate/util/__init__.py +9 -32
- brainstate/util/_others.py +1025 -0
- brainstate/util/_others_test.py +962 -0
- brainstate/util/_pretty_pytree.py +1301 -0
- brainstate/util/_pretty_pytree_test.py +675 -0
- brainstate/util/{pretty_repr.py → _pretty_repr.py} +161 -27
- brainstate/util/_pretty_repr_test.py +696 -0
- brainstate/util/filter.py +557 -81
- brainstate/util/filter_test.py +912 -0
- brainstate/util/struct.py +769 -382
- brainstate/util/struct_test.py +602 -0
- {brainstate-0.1.9.dist-info → brainstate-0.2.0.dist-info}/METADATA +34 -17
- brainstate-0.2.0.dist-info/RECORD +111 -0
- brainstate/augment/__init__.py +0 -30
- brainstate/augment/_eval_shape.py +0 -99
- brainstate/augment/_mapping.py +0 -1060
- brainstate/augment/_mapping_test.py +0 -597
- brainstate/compile/__init__.py +0 -38
- brainstate/compile/_ad_checkpoint.py +0 -204
- brainstate/compile/_conditions.py +0 -256
- brainstate/compile/_make_jaxpr.py +0 -888
- brainstate/compile/_make_jaxpr_test.py +0 -156
- brainstate/compile/_util.py +0 -147
- brainstate/functional/__init__.py +0 -27
- brainstate/graph/_graph_node.py +0 -244
- brainstate/graph/_graph_node_test.py +0 -73
- brainstate/graph/_graph_operation_test.py +0 -563
- brainstate/init/__init__.py +0 -26
- brainstate/init/_base.py +0 -52
- brainstate/init/_generic.py +0 -244
- brainstate/init/_regular_inits.py +0 -105
- brainstate/init/_regular_inits_test.py +0 -50
- brainstate/nn/_inputs.py +0 -608
- brainstate/nn/_ltp.py +0 -28
- brainstate/nn/_neuron.py +0 -705
- brainstate/nn/_neuron_test.py +0 -161
- brainstate/nn/_others.py +0 -46
- brainstate/nn/_projection.py +0 -486
- brainstate/nn/_rate_rnns_test.py +0 -63
- brainstate/nn/_readout.py +0 -209
- brainstate/nn/_readout_test.py +0 -53
- brainstate/nn/_stp.py +0 -236
- brainstate/nn/_synapse.py +0 -505
- brainstate/nn/_synapse_test.py +0 -131
- brainstate/nn/_synaptic_projection.py +0 -423
- brainstate/nn/_synouts.py +0 -162
- brainstate/nn/_synouts_test.py +0 -57
- brainstate/nn/metrics.py +0 -388
- brainstate/optim/__init__.py +0 -38
- brainstate/optim/_base.py +0 -64
- brainstate/optim/_lr_scheduler.py +0 -448
- brainstate/optim/_lr_scheduler_test.py +0 -50
- brainstate/optim/_optax_optimizer.py +0 -152
- brainstate/optim/_optax_optimizer_test.py +0 -53
- brainstate/optim/_sgd_optimizer.py +0 -1104
- brainstate/random/_random_for_unit.py +0 -52
- brainstate/surrogate.py +0 -1957
- brainstate/transform.py +0 -23
- brainstate/util/caller.py +0 -98
- brainstate/util/others.py +0 -540
- brainstate/util/pretty_pytree.py +0 -945
- brainstate/util/pretty_pytree_test.py +0 -159
- brainstate/util/pretty_table.py +0 -2954
- brainstate/util/scaling.py +0 -258
- brainstate-0.1.9.dist-info/RECORD +0 -130
- {brainstate-0.1.9.dist-info → brainstate-0.2.0.dist-info}/WHEEL +0 -0
- {brainstate-0.1.9.dist-info → brainstate-0.2.0.dist-info}/licenses/LICENSE +0 -0
- {brainstate-0.1.9.dist-info → brainstate-0.2.0.dist-info}/top_level.txt +0 -0
brainstate/nn/_dropout_test.py
CHANGED
@@ -1,4 +1,4 @@
|
|
1
|
-
# Copyright 2024
|
1
|
+
# Copyright 2024 BrainX Ecosystem Limited. All Rights Reserved.
|
2
2
|
#
|
3
3
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
4
4
|
# you may not use this file except in compliance with the License.
|
@@ -14,24 +14,21 @@
|
|
14
14
|
# ==============================================================================
|
15
15
|
|
16
16
|
|
17
|
-
import unittest
|
18
|
-
|
19
17
|
import numpy as np
|
18
|
+
from absl.testing import absltest
|
19
|
+
from absl.testing import parameterized
|
20
20
|
|
21
21
|
import brainstate
|
22
22
|
|
23
23
|
|
24
|
-
class TestDropout(
|
24
|
+
class TestDropout(parameterized.TestCase):
|
25
25
|
|
26
|
-
def
|
27
|
-
|
26
|
+
def test_dropout_basic(self):
|
27
|
+
"""Test basic dropout functionality."""
|
28
28
|
dropout_layer = brainstate.nn.Dropout(0.5)
|
29
|
-
|
30
|
-
# Input data
|
31
|
-
input_data = np.arange(20)
|
29
|
+
input_data = np.arange(20, dtype=np.float32)
|
32
30
|
|
33
31
|
with brainstate.environ.context(fit=True):
|
34
|
-
# Apply dropout
|
35
32
|
output_data = dropout_layer(input_data)
|
36
33
|
|
37
34
|
# Check that the output has the same shape as the input
|
@@ -46,55 +43,435 @@ class TestDropout(unittest.TestCase):
|
|
46
43
|
expected_non_zero_elements = input_data[output_data != 0] * scale_factor
|
47
44
|
np.testing.assert_almost_equal(non_zero_elements, expected_non_zero_elements)
|
48
45
|
|
49
|
-
def
|
50
|
-
|
51
|
-
dropout_layer.
|
52
|
-
input_data = np.
|
53
|
-
|
54
|
-
|
55
|
-
|
56
|
-
|
57
|
-
|
58
|
-
|
59
|
-
|
60
|
-
|
61
|
-
|
62
|
-
|
63
|
-
|
64
|
-
|
65
|
-
|
66
|
-
|
67
|
-
|
68
|
-
|
69
|
-
|
70
|
-
|
71
|
-
|
72
|
-
|
73
|
-
|
74
|
-
|
46
|
+
def test_dropout_eval_mode(self):
|
47
|
+
"""Test that dropout is disabled in evaluation mode."""
|
48
|
+
dropout_layer = brainstate.nn.Dropout(0.5)
|
49
|
+
input_data = np.arange(20, dtype=np.float32)
|
50
|
+
|
51
|
+
with brainstate.environ.context(fit=False):
|
52
|
+
# Without fit context, dropout should be disabled
|
53
|
+
output_data = dropout_layer(input_data)
|
54
|
+
np.testing.assert_array_equal(input_data, output_data)
|
55
|
+
|
56
|
+
@parameterized.parameters(0.0, 0.2, 0.5, 0.8, 1.0)
|
57
|
+
def test_dropout_various_probs(self, prob):
|
58
|
+
"""Test dropout with various probabilities."""
|
59
|
+
dropout_layer = brainstate.nn.Dropout(prob)
|
60
|
+
input_data = brainstate.random.randn(1000)
|
61
|
+
|
62
|
+
with brainstate.environ.context(fit=True):
|
63
|
+
output_data = dropout_layer(input_data)
|
64
|
+
self.assertEqual(input_data.shape, output_data.shape)
|
65
|
+
|
66
|
+
if prob == 1.0:
|
67
|
+
# All elements should be kept (no dropout)
|
68
|
+
np.testing.assert_array_equal(input_data, output_data)
|
69
|
+
elif prob == 0.0:
|
70
|
+
# All elements should be dropped
|
71
|
+
np.testing.assert_array_equal(np.zeros_like(input_data), output_data)
|
72
|
+
|
73
|
+
def test_dropout_broadcast_dims(self):
|
74
|
+
"""Test dropout with broadcast dimensions."""
|
75
|
+
dropout_layer = brainstate.nn.Dropout(0.5, broadcast_dims=(1, 2))
|
76
|
+
input_data = brainstate.random.randn(4, 5, 6)
|
77
|
+
|
78
|
+
with brainstate.environ.context(fit=True):
|
79
|
+
output_data = dropout_layer(input_data)
|
80
|
+
self.assertEqual(input_data.shape, output_data.shape)
|
81
|
+
|
82
|
+
# Check that the same mask is applied across broadcast dimensions
|
83
|
+
for i in range(4):
|
84
|
+
channel_output = output_data[i]
|
85
|
+
# All elements in a channel should be either all zero or all scaled
|
86
|
+
is_zero = (channel_output == 0).all() or (channel_output != 0).all()
|
87
|
+
# Note: This might not always be true due to randomness, but generally should hold
|
88
|
+
|
89
|
+
def test_dropout_multidimensional(self):
|
90
|
+
"""Test dropout with multidimensional input."""
|
91
|
+
dropout_layer = brainstate.nn.Dropout(0.5)
|
92
|
+
input_data = brainstate.random.randn(10, 20, 30)
|
93
|
+
|
94
|
+
with brainstate.environ.context(fit=True):
|
95
|
+
output_data = dropout_layer(input_data)
|
96
|
+
self.assertEqual(input_data.shape, output_data.shape)
|
97
|
+
|
98
|
+
|
99
|
+
class TestDropout1d(parameterized.TestCase):
|
100
|
+
|
101
|
+
def test_dropout1d_basic(self):
|
102
|
+
"""Test basic Dropout1d functionality."""
|
103
|
+
dropout_layer = brainstate.nn.Dropout1d(prob=0.5)
|
104
|
+
input_data = brainstate.random.randn(2, 3, 4) # (N, C, L)
|
105
|
+
|
106
|
+
with brainstate.environ.context(fit=True):
|
107
|
+
output_data = dropout_layer(input_data)
|
108
|
+
self.assertEqual(input_data.shape, output_data.shape)
|
109
|
+
|
110
|
+
def test_dropout1d_channel_wise(self):
|
111
|
+
"""Test that Dropout1d applies dropout."""
|
112
|
+
dropout_layer = brainstate.nn.Dropout1d(prob=0.5)
|
113
|
+
input_data = brainstate.random.randn(2, 8, 10) # (N, C, L)
|
114
|
+
|
115
|
+
with brainstate.environ.context(fit=True):
|
116
|
+
output_data = dropout_layer(input_data)
|
117
|
+
# Just verify that dropout is applied (shape preserved, some zeros present)
|
118
|
+
self.assertEqual(input_data.shape, output_data.shape)
|
119
|
+
|
120
|
+
def test_dropout1d_without_batch(self):
|
121
|
+
"""Test Dropout1d with unbatched input (C, L)."""
|
122
|
+
dropout_layer = brainstate.nn.Dropout1d(prob=0.5)
|
123
|
+
input_data = brainstate.random.randn(3, 4) # (C, L)
|
124
|
+
|
125
|
+
with brainstate.environ.context(fit=True):
|
126
|
+
output_data = dropout_layer(input_data)
|
127
|
+
self.assertEqual(input_data.shape, output_data.shape)
|
128
|
+
|
129
|
+
def test_dropout1d_eval_mode(self):
|
130
|
+
"""Test that Dropout1d is disabled in eval mode."""
|
131
|
+
dropout_layer = brainstate.nn.Dropout1d(prob=0.5)
|
132
|
+
input_data = brainstate.random.randn(2, 3, 4)
|
133
|
+
|
134
|
+
with brainstate.environ.context(fit=False):
|
135
|
+
output_data = dropout_layer(input_data)
|
136
|
+
np.testing.assert_array_equal(input_data, output_data)
|
137
|
+
|
138
|
+
def test_dropout1d_invalid_shape(self):
|
139
|
+
"""Test that Dropout1d raises error for invalid input shape."""
|
140
|
+
dropout_layer = brainstate.nn.Dropout1d(prob=0.5)
|
141
|
+
input_data = brainstate.random.randn(2, 3, 4, 5) # 4D input is invalid for Dropout1d
|
142
|
+
|
143
|
+
with brainstate.environ.context(fit=True):
|
144
|
+
with self.assertRaises(RuntimeError):
|
145
|
+
dropout_layer(input_data)
|
146
|
+
|
147
|
+
|
148
|
+
class TestDropout2d(parameterized.TestCase):
|
149
|
+
|
150
|
+
def test_dropout2d_basic(self):
|
151
|
+
"""Test basic Dropout2d functionality."""
|
152
|
+
dropout_layer = brainstate.nn.Dropout2d(prob=0.5)
|
153
|
+
input_data = brainstate.random.randn(2, 3, 4, 5) # (N, C, H, W)
|
154
|
+
|
155
|
+
with brainstate.environ.context(fit=True):
|
156
|
+
output_data = dropout_layer(input_data)
|
157
|
+
self.assertEqual(input_data.shape, output_data.shape)
|
158
|
+
self.assertTrue(np.any(output_data == 0))
|
159
|
+
|
160
|
+
def test_dropout2d_channel_wise(self):
|
161
|
+
"""Test that Dropout2d applies dropout."""
|
162
|
+
dropout_layer = brainstate.nn.Dropout2d(prob=0.5)
|
163
|
+
input_data = brainstate.random.randn(2, 8, 4, 5) # (N, C, H, W)
|
164
|
+
|
165
|
+
with brainstate.environ.context(fit=True):
|
166
|
+
output_data = dropout_layer(input_data)
|
167
|
+
# Just verify that dropout is applied (shape preserved, some zeros present)
|
168
|
+
self.assertEqual(input_data.shape, output_data.shape)
|
169
|
+
|
170
|
+
def test_dropout2d_without_batch(self):
|
171
|
+
"""Test Dropout2d with unbatched input (C, H, W)."""
|
172
|
+
dropout_layer = brainstate.nn.Dropout2d(prob=0.5)
|
173
|
+
input_data = brainstate.random.randn(3, 4, 5) # (C, H, W)
|
174
|
+
|
175
|
+
with brainstate.environ.context(fit=True):
|
176
|
+
output_data = dropout_layer(input_data)
|
177
|
+
self.assertEqual(input_data.shape, output_data.shape)
|
178
|
+
|
179
|
+
def test_dropout2d_scaling(self):
|
180
|
+
"""Test that Dropout2d correctly scales non-dropped elements."""
|
75
181
|
dropout_layer = brainstate.nn.Dropout2d(prob=0.5)
|
76
|
-
input_data =
|
182
|
+
input_data = brainstate.random.randn(2, 3, 4, 5)
|
183
|
+
|
184
|
+
with brainstate.environ.context(fit=True):
|
185
|
+
output_data = dropout_layer(input_data)
|
186
|
+
scale_factor = 1 / (1 - 0.5)
|
187
|
+
mask = ~np.isclose(output_data, 0)
|
188
|
+
non_zero_elements = output_data[mask]
|
189
|
+
expected_non_zero_elements = input_data[mask] * scale_factor
|
190
|
+
np.testing.assert_allclose(non_zero_elements, expected_non_zero_elements, rtol=1e-5)
|
191
|
+
|
192
|
+
def test_dropout2d_eval_mode(self):
|
193
|
+
"""Test that Dropout2d is disabled in eval mode."""
|
194
|
+
dropout_layer = brainstate.nn.Dropout2d(prob=0.5)
|
195
|
+
input_data = brainstate.random.randn(2, 3, 4, 5)
|
196
|
+
|
197
|
+
with brainstate.environ.context(fit=False):
|
198
|
+
output_data = dropout_layer(input_data)
|
199
|
+
np.testing.assert_array_equal(input_data, output_data)
|
200
|
+
|
201
|
+
|
202
|
+
class TestDropout3d(parameterized.TestCase):
|
203
|
+
|
204
|
+
def test_dropout3d_basic(self):
|
205
|
+
"""Test basic Dropout3d functionality."""
|
206
|
+
dropout_layer = brainstate.nn.Dropout3d(prob=0.5)
|
207
|
+
input_data = brainstate.random.randn(2, 3, 4, 5, 6) # (N, C, D, H, W)
|
208
|
+
|
77
209
|
with brainstate.environ.context(fit=True):
|
78
210
|
output_data = dropout_layer(input_data)
|
79
|
-
|
80
|
-
|
81
|
-
scale_factor = 1 / (1 - 0.5)
|
82
|
-
non_zero_elements = output_data[output_data != 0]
|
83
|
-
expected_non_zero_elements = input_data[output_data != 0] * scale_factor
|
84
|
-
np.testing.assert_almost_equal(non_zero_elements, expected_non_zero_elements, decimal=4)
|
211
|
+
self.assertEqual(input_data.shape, output_data.shape)
|
212
|
+
self.assertTrue(np.any(output_data == 0))
|
85
213
|
|
86
|
-
def
|
214
|
+
def test_dropout3d_channel_wise(self):
|
215
|
+
"""Test that Dropout3d applies dropout."""
|
87
216
|
dropout_layer = brainstate.nn.Dropout3d(prob=0.5)
|
88
|
-
input_data =
|
217
|
+
input_data = brainstate.random.randn(2, 8, 4, 5, 6) # (N, C, D, H, W)
|
218
|
+
|
89
219
|
with brainstate.environ.context(fit=True):
|
90
220
|
output_data = dropout_layer(input_data)
|
91
|
-
|
92
|
-
|
93
|
-
|
94
|
-
|
95
|
-
|
96
|
-
|
221
|
+
# Just verify that dropout is applied (shape preserved, some zeros present)
|
222
|
+
self.assertEqual(input_data.shape, output_data.shape)
|
223
|
+
|
224
|
+
def test_dropout3d_without_batch(self):
|
225
|
+
"""Test Dropout3d with unbatched input (C, D, H, W)."""
|
226
|
+
dropout_layer = brainstate.nn.Dropout3d(prob=0.5)
|
227
|
+
input_data = brainstate.random.randn(3, 4, 5, 6) # (C, D, H, W)
|
228
|
+
|
229
|
+
with brainstate.environ.context(fit=True):
|
230
|
+
output_data = dropout_layer(input_data)
|
231
|
+
self.assertEqual(input_data.shape, output_data.shape)
|
232
|
+
|
233
|
+
def test_dropout3d_scaling(self):
|
234
|
+
"""Test that Dropout3d correctly scales non-dropped elements."""
|
235
|
+
dropout_layer = brainstate.nn.Dropout3d(prob=0.5)
|
236
|
+
input_data = brainstate.random.randn(2, 3, 4, 5, 6)
|
237
|
+
|
238
|
+
with brainstate.environ.context(fit=True):
|
239
|
+
output_data = dropout_layer(input_data)
|
240
|
+
scale_factor = 1 / (1 - 0.5)
|
241
|
+
mask = ~np.isclose(output_data, 0)
|
242
|
+
non_zero_elements = output_data[mask]
|
243
|
+
expected_non_zero_elements = input_data[mask] * scale_factor
|
244
|
+
np.testing.assert_allclose(non_zero_elements, expected_non_zero_elements, rtol=1e-5)
|
245
|
+
|
246
|
+
def test_dropout3d_eval_mode(self):
|
247
|
+
"""Test that Dropout3d is disabled in eval mode."""
|
248
|
+
dropout_layer = brainstate.nn.Dropout3d(prob=0.5)
|
249
|
+
input_data = brainstate.random.randn(2, 3, 4, 5, 6)
|
250
|
+
|
251
|
+
with brainstate.environ.context(fit=False):
|
252
|
+
output_data = dropout_layer(input_data)
|
253
|
+
np.testing.assert_array_equal(input_data, output_data)
|
254
|
+
|
255
|
+
|
256
|
+
class TestAlphaDropout(parameterized.TestCase):
|
257
|
+
|
258
|
+
def test_alphadropout_basic(self):
|
259
|
+
"""Test basic AlphaDropout functionality."""
|
260
|
+
dropout_layer = brainstate.nn.AlphaDropout(prob=0.5)
|
261
|
+
input_data = brainstate.random.randn(100, 50)
|
262
|
+
|
263
|
+
with brainstate.environ.context(fit=True):
|
264
|
+
output_data = dropout_layer(input_data)
|
265
|
+
self.assertEqual(input_data.shape, output_data.shape)
|
266
|
+
|
267
|
+
def test_alphadropout_self_normalizing(self):
|
268
|
+
"""Test that AlphaDropout maintains zero mean and unit variance."""
|
269
|
+
dropout_layer = brainstate.nn.AlphaDropout(prob=0.5)
|
270
|
+
# Create input with zero mean and unit variance
|
271
|
+
input_data = brainstate.random.randn(10000)
|
272
|
+
|
273
|
+
with brainstate.environ.context(fit=True):
|
274
|
+
output_data = dropout_layer(input_data)
|
275
|
+
|
276
|
+
# The output should approximately maintain zero mean and unit variance
|
277
|
+
output_mean = np.mean(output_data)
|
278
|
+
output_std = np.std(output_data)
|
279
|
+
|
280
|
+
# Allow some tolerance due to randomness
|
281
|
+
self.assertAlmostEqual(output_mean, 0.0, delta=0.1)
|
282
|
+
self.assertAlmostEqual(output_std, 1.0, delta=0.2)
|
283
|
+
|
284
|
+
def test_alphadropout_alpha_value(self):
|
285
|
+
"""Test that dropped values are set to alpha (not zero)."""
|
286
|
+
dropout_layer = brainstate.nn.AlphaDropout(prob=0.5)
|
287
|
+
input_data = brainstate.random.randn(1000) + 5.0 # Shift to avoid confusion with alpha
|
288
|
+
|
289
|
+
with brainstate.environ.context(fit=True):
|
290
|
+
output_data = dropout_layer(input_data)
|
291
|
+
|
292
|
+
# AlphaDropout should not have zeros, but values close to transformed alpha
|
293
|
+
# After affine transformation: alpha * a + b
|
294
|
+
expected_dropped_value = dropout_layer.alpha * dropout_layer.a + dropout_layer.b
|
295
|
+
|
296
|
+
# Check that we have some values close to the expected dropped value
|
297
|
+
# (within reasonable tolerance due to numerical precision)
|
298
|
+
unique_vals = np.unique(np.round(output_data, 3))
|
299
|
+
self.assertGreater(len(unique_vals), 1) # Should have both dropped and kept values
|
300
|
+
|
301
|
+
def test_alphadropout_eval_mode(self):
|
302
|
+
"""Test that AlphaDropout is disabled in eval mode."""
|
303
|
+
dropout_layer = brainstate.nn.AlphaDropout(prob=0.5)
|
304
|
+
input_data = brainstate.random.randn(20, 16)
|
305
|
+
|
306
|
+
with brainstate.environ.context(fit=False):
|
307
|
+
output_data = dropout_layer(input_data)
|
308
|
+
np.testing.assert_array_equal(input_data, output_data)
|
309
|
+
|
310
|
+
@parameterized.parameters(0.2, 0.5, 0.8)
|
311
|
+
def test_alphadropout_various_probs(self, prob):
|
312
|
+
"""Test AlphaDropout with various probabilities."""
|
313
|
+
dropout_layer = brainstate.nn.AlphaDropout(prob=prob)
|
314
|
+
input_data = brainstate.random.randn(1000)
|
315
|
+
|
316
|
+
with brainstate.environ.context(fit=True):
|
317
|
+
output_data = dropout_layer(input_data)
|
318
|
+
self.assertEqual(input_data.shape, output_data.shape)
|
319
|
+
|
320
|
+
def test_alphadropout_multidimensional(self):
|
321
|
+
"""Test AlphaDropout with multidimensional input."""
|
322
|
+
dropout_layer = brainstate.nn.AlphaDropout(prob=0.5)
|
323
|
+
input_data = brainstate.random.randn(10, 20, 30)
|
324
|
+
|
325
|
+
with brainstate.environ.context(fit=True):
|
326
|
+
output_data = dropout_layer(input_data)
|
327
|
+
self.assertEqual(input_data.shape, output_data.shape)
|
328
|
+
|
329
|
+
|
330
|
+
class TestFeatureAlphaDropout(parameterized.TestCase):
|
331
|
+
|
332
|
+
def test_featurealphadropout_basic(self):
|
333
|
+
"""Test basic FeatureAlphaDropout functionality."""
|
334
|
+
dropout_layer = brainstate.nn.FeatureAlphaDropout(prob=0.5)
|
335
|
+
input_data = brainstate.random.randn(2, 16, 4, 32, 32) # (N, C, D, H, W)
|
336
|
+
|
337
|
+
with brainstate.environ.context(fit=True):
|
338
|
+
output_data = dropout_layer(input_data)
|
339
|
+
self.assertEqual(input_data.shape, output_data.shape)
|
340
|
+
|
341
|
+
def test_featurealphadropout_channel_wise(self):
|
342
|
+
"""Test that FeatureAlphaDropout drops entire channels."""
|
343
|
+
dropout_layer = brainstate.nn.FeatureAlphaDropout(prob=0.5, channel_axis=1)
|
344
|
+
input_data = brainstate.random.randn(2, 8, 10, 10) # (N, C, H, W)
|
345
|
+
|
346
|
+
with brainstate.environ.context(fit=True):
|
347
|
+
output_data = dropout_layer(input_data)
|
348
|
+
|
349
|
+
# Check that entire channels share the same dropout mask
|
350
|
+
for batch in range(2):
|
351
|
+
for channel in range(8):
|
352
|
+
channel_data = output_data[batch, channel, :, :]
|
353
|
+
# All elements in a channel should be either all from input or all alpha-transformed
|
354
|
+
unique_vals = np.unique(np.round(channel_data, 4))
|
355
|
+
# Due to the alpha transformation, we can't easily check, but shape should match
|
356
|
+
self.assertEqual(channel_data.shape, (10, 10))
|
357
|
+
|
358
|
+
def test_featurealphadropout_channel_axis(self):
|
359
|
+
"""Test FeatureAlphaDropout with different channel axes."""
|
360
|
+
# Test with channel_axis=-1
|
361
|
+
dropout_layer = brainstate.nn.FeatureAlphaDropout(prob=0.5, channel_axis=-1)
|
362
|
+
input_data = brainstate.random.randn(2, 10, 10, 8) # (N, H, W, C)
|
363
|
+
|
364
|
+
with brainstate.environ.context(fit=True):
|
365
|
+
output_data = dropout_layer(input_data)
|
366
|
+
self.assertEqual(input_data.shape, output_data.shape)
|
367
|
+
|
368
|
+
def test_featurealphadropout_eval_mode(self):
|
369
|
+
"""Test that FeatureAlphaDropout is disabled in eval mode."""
|
370
|
+
dropout_layer = brainstate.nn.FeatureAlphaDropout(prob=0.5)
|
371
|
+
input_data = brainstate.random.randn(2, 16, 4, 32, 32)
|
372
|
+
|
373
|
+
with brainstate.environ.context(fit=False):
|
374
|
+
output_data = dropout_layer(input_data)
|
375
|
+
np.testing.assert_array_equal(input_data, output_data)
|
376
|
+
|
377
|
+
def test_featurealphadropout_self_normalizing(self):
|
378
|
+
"""Test that FeatureAlphaDropout maintains self-normalizing properties."""
|
379
|
+
dropout_layer = brainstate.nn.FeatureAlphaDropout(prob=0.5, channel_axis=1)
|
380
|
+
# Create input with zero mean and unit variance
|
381
|
+
input_data = brainstate.random.randn(10, 32, 20)
|
382
|
+
|
383
|
+
with brainstate.environ.context(fit=True):
|
384
|
+
output_data = dropout_layer(input_data)
|
385
|
+
|
386
|
+
# The output should approximately maintain zero mean and unit variance
|
387
|
+
output_mean = np.mean(output_data)
|
388
|
+
output_std = np.std(output_data)
|
389
|
+
|
390
|
+
# Allow some tolerance due to randomness
|
391
|
+
self.assertAlmostEqual(output_mean, 0.0, delta=0.2)
|
392
|
+
self.assertAlmostEqual(output_std, 1.0, delta=0.3)
|
393
|
+
|
394
|
+
|
395
|
+
class TestDropoutFixed(parameterized.TestCase):
|
396
|
+
|
397
|
+
def test_dropoutfixed_basic(self):
|
398
|
+
"""Test basic DropoutFixed functionality."""
|
399
|
+
with brainstate.random.seed_context(42):
|
400
|
+
dropout_layer = brainstate.nn.DropoutFixed(in_size=(2, 3), prob=0.5)
|
401
|
+
dropout_layer.init_state(batch_size=2)
|
402
|
+
input_data = np.random.randn(2, 2, 3)
|
403
|
+
|
404
|
+
with brainstate.environ.context(fit=True):
|
405
|
+
output_data = dropout_layer.update(input_data)
|
406
|
+
self.assertEqual(input_data.shape, output_data.shape)
|
407
|
+
self.assertTrue(np.any(output_data == 0))
|
408
|
+
|
409
|
+
def test_dropoutfixed_mask_persistence(self):
|
410
|
+
"""Test that DropoutFixed uses the same mask across multiple calls."""
|
411
|
+
with brainstate.random.seed_context(42):
|
412
|
+
dropout_layer = brainstate.nn.DropoutFixed(in_size=(10,), prob=0.5)
|
413
|
+
dropout_layer.init_state(batch_size=5)
|
414
|
+
|
415
|
+
input_data1 = brainstate.random.randn(5, 10)
|
416
|
+
input_data2 = brainstate.random.randn(5, 10)
|
417
|
+
|
418
|
+
with brainstate.environ.context(fit=True):
|
419
|
+
output_data1 = dropout_layer.update(input_data1)
|
420
|
+
output_data2 = dropout_layer.update(input_data2)
|
421
|
+
|
422
|
+
# The dropout mask should be the same for both calls
|
423
|
+
mask1 = (output_data1 == 0)
|
424
|
+
mask2 = (output_data2 == 0)
|
425
|
+
np.testing.assert_array_equal(mask1, mask2)
|
426
|
+
|
427
|
+
def test_dropoutfixed_scaling(self):
|
428
|
+
"""Test that DropoutFixed correctly scales non-dropped elements."""
|
429
|
+
with brainstate.random.seed_context(42):
|
430
|
+
dropout_layer = brainstate.nn.DropoutFixed(in_size=(2, 3), prob=0.5)
|
431
|
+
dropout_layer.init_state(batch_size=2)
|
432
|
+
input_data = np.random.randn(2, 2, 3)
|
433
|
+
|
434
|
+
with brainstate.environ.context(fit=True):
|
435
|
+
output_data = dropout_layer.update(input_data)
|
436
|
+
scale_factor = 1 / (1 - 0.5)
|
437
|
+
non_zero_elements = output_data[output_data != 0]
|
438
|
+
expected_non_zero_elements = input_data[output_data != 0] * scale_factor
|
439
|
+
np.testing.assert_almost_equal(non_zero_elements, expected_non_zero_elements, decimal=3)
|
440
|
+
|
441
|
+
def test_dropoutfixed_eval_mode(self):
|
442
|
+
"""Test that DropoutFixed is disabled in eval mode."""
|
443
|
+
with brainstate.random.seed_context(42):
|
444
|
+
dropout_layer = brainstate.nn.DropoutFixed(in_size=(2, 3), prob=0.5)
|
445
|
+
dropout_layer.init_state(batch_size=2)
|
446
|
+
input_data = np.random.randn(2, 2, 3)
|
447
|
+
|
448
|
+
with brainstate.environ.context(fit=False):
|
449
|
+
output_data = dropout_layer.update(input_data)
|
450
|
+
np.testing.assert_array_equal(input_data, output_data)
|
451
|
+
|
452
|
+
def test_dropoutfixed_shape_mismatch(self):
|
453
|
+
"""Test that DropoutFixed raises error for shape mismatch."""
|
454
|
+
with brainstate.random.seed_context(42):
|
455
|
+
dropout_layer = brainstate.nn.DropoutFixed(in_size=(2, 3), prob=0.5)
|
456
|
+
dropout_layer.init_state(batch_size=2)
|
457
|
+
input_data = np.random.randn(3, 2, 3) # Wrong batch size
|
458
|
+
|
459
|
+
with brainstate.environ.context(fit=True):
|
460
|
+
with self.assertRaises(ValueError):
|
461
|
+
dropout_layer.update(input_data)
|
462
|
+
|
463
|
+
@parameterized.parameters(0.2, 0.5, 0.8)
|
464
|
+
def test_dropoutfixed_various_probs(self, prob):
|
465
|
+
"""Test DropoutFixed with various probabilities."""
|
466
|
+
with brainstate.random.seed_context(42):
|
467
|
+
dropout_layer = brainstate.nn.DropoutFixed(in_size=(10,), prob=prob)
|
468
|
+
dropout_layer.init_state(batch_size=5)
|
469
|
+
input_data = brainstate.random.randn(5, 10)
|
470
|
+
|
471
|
+
with brainstate.environ.context(fit=True):
|
472
|
+
output_data = dropout_layer.update(input_data)
|
473
|
+
self.assertEqual(input_data.shape, output_data.shape)
|
97
474
|
|
98
475
|
|
99
476
|
if __name__ == '__main__':
|
100
|
-
|
477
|
+
absltest.main()
|
brainstate/nn/_dynamics.py
CHANGED
@@ -1,4 +1,4 @@
|
|
1
|
-
# Copyright 2024
|
1
|
+
# Copyright 2024 BrainX Ecosystem Limited. All Rights Reserved.
|
2
2
|
#
|
3
3
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
4
4
|
# you may not use this file except in compliance with the License.
|
@@ -47,59 +47,18 @@ from ._delay import StateWithDelay, Delay
|
|
47
47
|
from ._module import Module
|
48
48
|
|
49
49
|
__all__ = [
|
50
|
-
'DynamicsGroup',
|
51
|
-
'
|
50
|
+
'DynamicsGroup',
|
51
|
+
'Dynamics',
|
52
|
+
'Prefetch',
|
53
|
+
'PrefetchDelay',
|
54
|
+
'PrefetchDelayAt',
|
55
|
+
'OutputDelayAt',
|
52
56
|
]
|
53
57
|
|
54
58
|
T = TypeVar('T')
|
55
59
|
_max_order = 10
|
56
60
|
|
57
61
|
|
58
|
-
class Projection(Module):
|
59
|
-
"""
|
60
|
-
Base class for synaptic projection modules in neural network modeling.
|
61
|
-
|
62
|
-
This class defines the interface for modules that handle projections between
|
63
|
-
neural populations. Projections process input signals and transform them
|
64
|
-
before they reach the target neurons, implementing the connectivity patterns
|
65
|
-
in neural networks.
|
66
|
-
|
67
|
-
In the BrainState execution order, Projection modules are updated before
|
68
|
-
Dynamics modules, following the natural information flow in neural systems:
|
69
|
-
1. Projections process inputs (synaptic transmission)
|
70
|
-
2. Dynamics update neuron states (neural integration)
|
71
|
-
|
72
|
-
The Projection class does not implement the update logic directly but delegates
|
73
|
-
to its child nodes. If no child nodes exist, it raises a ValueError.
|
74
|
-
|
75
|
-
Parameters
|
76
|
-
----------
|
77
|
-
*args : Any
|
78
|
-
Arguments passed to the parent Module class.
|
79
|
-
**kwargs : Any
|
80
|
-
Keyword arguments passed to the parent Module class.
|
81
|
-
|
82
|
-
Raises
|
83
|
-
------
|
84
|
-
ValueError
|
85
|
-
If the update() method is called but no child nodes are defined.
|
86
|
-
|
87
|
-
Notes
|
88
|
-
-----
|
89
|
-
Derived classes should implement specific projection behaviors, such as
|
90
|
-
dense connectivity, sparse connectivity, or specific weight update rules.
|
91
|
-
"""
|
92
|
-
__module__ = 'brainstate.nn'
|
93
|
-
|
94
|
-
def update(self, *args, **kwargs):
|
95
|
-
sub_nodes = tuple(self.nodes(allowed_hierarchy=(1, 1)).values())
|
96
|
-
if len(sub_nodes):
|
97
|
-
for node in sub_nodes:
|
98
|
-
node(*args, **kwargs)
|
99
|
-
else:
|
100
|
-
raise ValueError('Do not implement the update() function.')
|
101
|
-
|
102
|
-
|
103
62
|
class Dynamics(Module):
|
104
63
|
"""
|
105
64
|
Base class for implementing neural dynamics models in BrainState.
|
@@ -445,12 +404,12 @@ class Dynamics(Module):
|
|
445
404
|
----------
|
446
405
|
init : Any
|
447
406
|
The initial value to which all current inputs will be added.
|
448
|
-
*args
|
407
|
+
*args
|
449
408
|
Variable length argument list passed to each current input function.
|
450
409
|
label : Optional[str], default=None
|
451
410
|
If provided, only process current inputs with this label prefix.
|
452
411
|
When None, process all current inputs regardless of label.
|
453
|
-
**kwargs
|
412
|
+
**kwargs
|
454
413
|
Arbitrary keyword arguments passed to each current input function.
|
455
414
|
|
456
415
|
Returns
|
@@ -481,7 +440,7 @@ class Dynamics(Module):
|
|
481
440
|
init = init + out(*args, **kwargs)
|
482
441
|
except Exception as e:
|
483
442
|
raise ValueError(
|
484
|
-
f'Error in
|
443
|
+
f'Error in current input value {key}: {out}\n'
|
485
444
|
f'Error: {e}'
|
486
445
|
) from e
|
487
446
|
else:
|
@@ -489,7 +448,7 @@ class Dynamics(Module):
|
|
489
448
|
init = init + out
|
490
449
|
except Exception as e:
|
491
450
|
raise ValueError(
|
492
|
-
f'Error in
|
451
|
+
f'Error in current input value {key}: {out}\n'
|
493
452
|
f'Error: {e}'
|
494
453
|
) from e
|
495
454
|
self._current_inputs.pop(key)
|
@@ -513,12 +472,12 @@ class Dynamics(Module):
|
|
513
472
|
----------
|
514
473
|
init : Any
|
515
474
|
The initial value to which all delta inputs will be added.
|
516
|
-
*args
|
475
|
+
*args
|
517
476
|
Variable length argument list passed to each delta input function.
|
518
477
|
label : Optional[str], default=None
|
519
478
|
If provided, only process delta inputs with this label prefix.
|
520
479
|
When None, process all delta inputs regardless of label.
|
521
|
-
**kwargs
|
480
|
+
**kwargs
|
522
481
|
Arbitrary keyword arguments passed to each delta input function.
|
523
482
|
|
524
483
|
Returns
|
@@ -1228,42 +1187,7 @@ def maybe_init_prefetch(target, *args, **kwargs):
|
|
1228
1187
|
# delay.register_delay(*target.delay_time)
|
1229
1188
|
|
1230
1189
|
|
1231
|
-
|
1232
|
-
"""
|
1233
|
-
A group of :py:class:`~.Module` in which the updating order does not matter.
|
1234
|
-
|
1235
|
-
Args:
|
1236
|
-
children_as_tuple: The children objects.
|
1237
|
-
children_as_dict: The children objects.
|
1238
|
-
"""
|
1239
|
-
|
1240
|
-
__module__ = 'brainstate.nn'
|
1241
|
-
|
1242
|
-
if not TYPE_CHECKING:
|
1243
|
-
def __init__(self, *children_as_tuple, **children_as_dict):
|
1244
|
-
super().__init__()
|
1245
|
-
self.layers_tuple = tuple(children_as_tuple)
|
1246
|
-
self.layers_dict = dict(children_as_dict)
|
1247
|
-
|
1248
|
-
def update(self, *args, **kwargs):
|
1249
|
-
"""
|
1250
|
-
Update function of a network.
|
1251
|
-
|
1252
|
-
In this update function, the update functions in children systems are iteratively called.
|
1253
|
-
"""
|
1254
|
-
projs, dyns, others = self.nodes(allowed_hierarchy=(1, 1)).split(Projection, Dynamics)
|
1255
|
-
|
1256
|
-
# update nodes of projections
|
1257
|
-
for node in projs.values():
|
1258
|
-
node()
|
1259
|
-
|
1260
|
-
# update nodes of dynamics
|
1261
|
-
for node in dyns.values():
|
1262
|
-
node()
|
1263
|
-
|
1264
|
-
# update nodes with other types, including delays, ...
|
1265
|
-
for node in others.values():
|
1266
|
-
node()
|
1190
|
+
DynamicsGroup = Module
|
1267
1191
|
|
1268
1192
|
|
1269
1193
|
def receive_update_output(cls: object):
|