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/init/_generic.py
DELETED
@@ -1,244 +0,0 @@
|
|
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 Union, Callable, Optional, Sequence
|
19
|
-
|
20
|
-
import brainunit as bu
|
21
|
-
import jax
|
22
|
-
import numpy as np
|
23
|
-
|
24
|
-
from brainstate._state import State
|
25
|
-
from brainstate._utils import set_module_as
|
26
|
-
from brainstate.typing import ArrayLike
|
27
|
-
from ._base import to_size
|
28
|
-
|
29
|
-
__all__ = [
|
30
|
-
'param',
|
31
|
-
'state',
|
32
|
-
'noise',
|
33
|
-
]
|
34
|
-
|
35
|
-
|
36
|
-
def _is_scalar(x):
|
37
|
-
return bu.math.isscalar(x)
|
38
|
-
|
39
|
-
|
40
|
-
def are_broadcastable_shapes(shape1, shape2):
|
41
|
-
"""
|
42
|
-
Check if two shapes are broadcastable.
|
43
|
-
|
44
|
-
Parameters:
|
45
|
-
- shape1: Tuple[int], the shape of the first array.
|
46
|
-
- shape2: Tuple[int], the shape of the second array.
|
47
|
-
|
48
|
-
Returns:
|
49
|
-
- bool: True if shapes are broadcastable, False otherwise.
|
50
|
-
"""
|
51
|
-
# Reverse the shapes to compare from the last dimension
|
52
|
-
shape1_reversed = shape1[::-1]
|
53
|
-
shape2_reversed = shape2[::-1]
|
54
|
-
|
55
|
-
# Iterate over the dimensions of the shorter shape
|
56
|
-
for dim1, dim2 in zip(shape1_reversed, shape2_reversed):
|
57
|
-
# Check if the dimensions are not equal and neither is 1
|
58
|
-
if dim1 != dim2 and 1 not in (dim1, dim2):
|
59
|
-
return False
|
60
|
-
|
61
|
-
# If all dimensions are compatible, the shapes are broadcastable
|
62
|
-
return True
|
63
|
-
|
64
|
-
|
65
|
-
def _expand_params_to_match_sizes(params, sizes):
|
66
|
-
"""
|
67
|
-
Expand the dimensions of params to match the dimensions of sizes.
|
68
|
-
|
69
|
-
Parameters:
|
70
|
-
- params: jax.Array or np.ndarray, the parameter array to be expanded.
|
71
|
-
- sizes: tuple[int] or list[int], the target shape dimensions.
|
72
|
-
|
73
|
-
Returns:
|
74
|
-
- Expanded params with dimensions matching sizes.
|
75
|
-
"""
|
76
|
-
params_dim = params.ndim
|
77
|
-
sizes_dim = len(sizes)
|
78
|
-
dim_diff = sizes_dim - params_dim
|
79
|
-
|
80
|
-
# Add new axes to params if it has fewer dimensions than sizes
|
81
|
-
for _ in range(dim_diff):
|
82
|
-
params = bu.math.expand_dims(params, axis=0) # Add new axis at the last dimension
|
83
|
-
return params
|
84
|
-
|
85
|
-
|
86
|
-
@set_module_as('brainstate.init')
|
87
|
-
def param(
|
88
|
-
parameter: Union[Callable, ArrayLike, State],
|
89
|
-
sizes: Union[int, Sequence[int]],
|
90
|
-
batch_size: Optional[int] = None,
|
91
|
-
allow_none: bool = True,
|
92
|
-
allow_scalar: bool = True,
|
93
|
-
):
|
94
|
-
"""Initialize parameters.
|
95
|
-
|
96
|
-
Parameters
|
97
|
-
----------
|
98
|
-
parameter: callable, ArrayLike, State
|
99
|
-
The initialization of the parameter.
|
100
|
-
- If it is None, the created parameter will be None.
|
101
|
-
- If it is a callable function :math:`f`, the ``f(size)`` will be returned.
|
102
|
-
- If it is an instance of :py:class:`init.Initializer``, the ``f(size)`` will be returned.
|
103
|
-
- If it is a tensor, then this function check whether ``tensor.shape`` is equal to the given ``size``.
|
104
|
-
sizes: int, sequence of int
|
105
|
-
The shape of the parameter.
|
106
|
-
batch_size: int
|
107
|
-
The batch size.
|
108
|
-
allow_none: bool
|
109
|
-
Whether allow the parameter is None.
|
110
|
-
allow_scalar: bool
|
111
|
-
Whether allow the parameter is a scalar value.
|
112
|
-
|
113
|
-
Returns
|
114
|
-
-------
|
115
|
-
param: ArrayType, float, int, bool, None
|
116
|
-
The initialized parameter.
|
117
|
-
|
118
|
-
See Also
|
119
|
-
--------
|
120
|
-
noise, state
|
121
|
-
"""
|
122
|
-
# Check if the parameter is None
|
123
|
-
if parameter is None:
|
124
|
-
if allow_none:
|
125
|
-
return None
|
126
|
-
else:
|
127
|
-
raise ValueError(f'Expect a parameter with type of float, ArrayType, Initializer, or '
|
128
|
-
f'Callable function, but we got None. ')
|
129
|
-
|
130
|
-
# Check if the parameter is a scalar value
|
131
|
-
if allow_scalar and _is_scalar(parameter):
|
132
|
-
return parameter
|
133
|
-
|
134
|
-
# Convert sizes to a tuple
|
135
|
-
sizes = tuple(to_size(sizes))
|
136
|
-
|
137
|
-
# Check if the parameter is a callable function
|
138
|
-
if callable(parameter):
|
139
|
-
if batch_size is not None:
|
140
|
-
sizes = (batch_size,) + sizes
|
141
|
-
return parameter(sizes)
|
142
|
-
elif isinstance(parameter, (np.ndarray, jax.Array, bu.Quantity, State)):
|
143
|
-
parameter = parameter
|
144
|
-
else:
|
145
|
-
raise ValueError(f'Unknown parameter type: {type(parameter)}')
|
146
|
-
|
147
|
-
# Check if the shape of the parameter matches the given size
|
148
|
-
if not are_broadcastable_shapes(parameter.shape, sizes):
|
149
|
-
raise ValueError(f'The shape of the parameter {parameter.shape} does not match with the given size {sizes}')
|
150
|
-
|
151
|
-
# Expand the parameter to match the given batch size
|
152
|
-
param_value = parameter.value if isinstance(parameter, State) else parameter
|
153
|
-
if batch_size is not None:
|
154
|
-
if param_value.ndim <= len(sizes):
|
155
|
-
# add a new axis to the params so that it matches the dimensionality of the given shape ``sizes``
|
156
|
-
param_value = _expand_params_to_match_sizes(param_value, sizes)
|
157
|
-
param_value = bu.math.repeat(
|
158
|
-
bu.math.expand_dims(param_value, axis=0),
|
159
|
-
batch_size,
|
160
|
-
axis=0
|
161
|
-
)
|
162
|
-
else:
|
163
|
-
if param_value.shape[0] != batch_size:
|
164
|
-
raise ValueError(f'The batch size of the parameter {param_value.shape[0]} '
|
165
|
-
f'does not match with the given batch size {batch_size}')
|
166
|
-
return type(parameter)(param_value) if isinstance(parameter, State) else param_value
|
167
|
-
|
168
|
-
|
169
|
-
@set_module_as('brainstate.init')
|
170
|
-
def state(
|
171
|
-
init: Union[Callable, jax.typing.ArrayLike],
|
172
|
-
sizes: Union[int, Sequence[int]] = None,
|
173
|
-
batch_size: Optional[int] = None,
|
174
|
-
):
|
175
|
-
"""
|
176
|
-
Initialize a :math:`~.State` from a callable function or a data.
|
177
|
-
"""
|
178
|
-
sizes = to_size(sizes)
|
179
|
-
if callable(init):
|
180
|
-
if sizes is None:
|
181
|
-
raise ValueError('"varshape" cannot be None when data is a callable function.')
|
182
|
-
sizes = list(sizes)
|
183
|
-
if isinstance(batch_size, int):
|
184
|
-
sizes.insert(0, batch_size)
|
185
|
-
return State(init(sizes))
|
186
|
-
|
187
|
-
else:
|
188
|
-
if sizes is not None:
|
189
|
-
if bu.math.shape(init) != sizes:
|
190
|
-
raise ValueError(f'The shape of "data" {bu.math.shape(init)} does not match with "var_shape" {sizes}')
|
191
|
-
if isinstance(batch_size, int):
|
192
|
-
batch_size = batch_size
|
193
|
-
data = State(
|
194
|
-
bu.math.repeat(
|
195
|
-
bu.math.expand_dims(init, axis=0),
|
196
|
-
batch_size,
|
197
|
-
axis=0
|
198
|
-
)
|
199
|
-
)
|
200
|
-
else:
|
201
|
-
data = State(init)
|
202
|
-
return data
|
203
|
-
|
204
|
-
|
205
|
-
@set_module_as('brainstate.init')
|
206
|
-
def noise(
|
207
|
-
noises: Optional[Union[ArrayLike, Callable]],
|
208
|
-
size: Union[int, Sequence[int]],
|
209
|
-
num_vars: int = 1,
|
210
|
-
noise_idx: int = 0,
|
211
|
-
) -> Optional[Callable]:
|
212
|
-
"""Initialize a noise function.
|
213
|
-
|
214
|
-
Parameters
|
215
|
-
----------
|
216
|
-
noises: Any
|
217
|
-
size: Shape
|
218
|
-
The size of the noise.
|
219
|
-
num_vars: int
|
220
|
-
The number of variables.
|
221
|
-
noise_idx: int
|
222
|
-
The index of the current noise among all noise variables.
|
223
|
-
|
224
|
-
Returns
|
225
|
-
-------
|
226
|
-
noise_func: function, None
|
227
|
-
The noise function.
|
228
|
-
|
229
|
-
See Also
|
230
|
-
--------
|
231
|
-
variable_, parameter, delay
|
232
|
-
|
233
|
-
"""
|
234
|
-
if callable(noises):
|
235
|
-
return noises
|
236
|
-
elif noises is None:
|
237
|
-
return None
|
238
|
-
else:
|
239
|
-
noises = param(noises, size, allow_none=False)
|
240
|
-
if num_vars > 1:
|
241
|
-
noises_ = [None] * num_vars
|
242
|
-
noises_[noise_idx] = noises
|
243
|
-
noises = tuple(noises_)
|
244
|
-
return lambda *args, **kwargs: noises
|
@@ -1,105 +0,0 @@
|
|
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
|
-
import brainunit as u
|
19
|
-
|
20
|
-
from brainstate import environ
|
21
|
-
from ._base import Initializer, to_size
|
22
|
-
|
23
|
-
__all__ = [
|
24
|
-
'ZeroInit',
|
25
|
-
'Constant',
|
26
|
-
'Identity',
|
27
|
-
]
|
28
|
-
|
29
|
-
|
30
|
-
class ZeroInit(Initializer):
|
31
|
-
"""Zero initializer.
|
32
|
-
|
33
|
-
Initialize the weights with zeros.
|
34
|
-
"""
|
35
|
-
__module__ = 'brainstate.init'
|
36
|
-
|
37
|
-
def __init__(self, unit: u.Unit = u.UNITLESS):
|
38
|
-
super(ZeroInit, self).__init__()
|
39
|
-
self.unit = unit
|
40
|
-
|
41
|
-
def __call__(self, shape, dtype=None):
|
42
|
-
dtype = dtype or environ.dftype()
|
43
|
-
shape = to_size(shape)
|
44
|
-
return u.maybe_decimal(u.math.zeros(shape, dtype=dtype, unit=self.unit))
|
45
|
-
|
46
|
-
|
47
|
-
class Constant(Initializer):
|
48
|
-
"""Constant initializer.
|
49
|
-
|
50
|
-
Initialize the weights with the given values.
|
51
|
-
|
52
|
-
Parameters
|
53
|
-
----------
|
54
|
-
value : float, int, bm.ndarray
|
55
|
-
The value to specify.
|
56
|
-
"""
|
57
|
-
__module__ = 'brainstate.init'
|
58
|
-
|
59
|
-
def __init__(self, value=1., ):
|
60
|
-
super(Constant, self).__init__()
|
61
|
-
self.value = value
|
62
|
-
|
63
|
-
def __call__(self, shape, dtype=None):
|
64
|
-
dtype = dtype or environ.dftype()
|
65
|
-
shape = to_size(shape)
|
66
|
-
return u.maybe_decimal(u.math.full(shape, self.value, dtype=dtype))
|
67
|
-
|
68
|
-
|
69
|
-
class Identity(Initializer):
|
70
|
-
"""Returns the identity matrix.
|
71
|
-
|
72
|
-
This initializer was proposed in (Le, et al., 2015) [1]_.
|
73
|
-
|
74
|
-
Parameters
|
75
|
-
----------
|
76
|
-
value : float
|
77
|
-
The optional scaling factor.
|
78
|
-
|
79
|
-
Returns
|
80
|
-
-------
|
81
|
-
shape: tuple of int
|
82
|
-
The weight shape/size.
|
83
|
-
|
84
|
-
References
|
85
|
-
----------
|
86
|
-
.. [1] Le, Quoc V., Navdeep Jaitly, and Geoffrey E. Hinton. "A simple way to
|
87
|
-
initialize recurrent networks of rectified linear units." arXiv preprint
|
88
|
-
arXiv:1504.00941 (2015).
|
89
|
-
"""
|
90
|
-
__module__ = 'brainstate.init'
|
91
|
-
|
92
|
-
def __init__(self, value=1., unit: u.Unit = u.UNITLESS):
|
93
|
-
super(Identity, self).__init__()
|
94
|
-
self.value = value
|
95
|
-
self.unit = unit
|
96
|
-
|
97
|
-
def __call__(self, shape, dtype=None):
|
98
|
-
dtype = dtype or environ.dftype()
|
99
|
-
shape = to_size(shape)
|
100
|
-
if isinstance(shape, (tuple, list)):
|
101
|
-
if len(shape) > 2:
|
102
|
-
raise ValueError(f'Only support initialize 2D weights for {self.__class__.__name__}.')
|
103
|
-
r = u.math.eye(*shape, dtype=dtype)
|
104
|
-
r = u.math.fill_diagonal(r, self.value)
|
105
|
-
return u.maybe_decimal(u.Quantity(r, unit=self.unit))
|
@@ -1,50 +0,0 @@
|
|
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
|
-
import unittest
|
19
|
-
|
20
|
-
import brainstate
|
21
|
-
|
22
|
-
|
23
|
-
class TestZeroInit(unittest.TestCase):
|
24
|
-
def test_zero_init(self):
|
25
|
-
init = brainstate.init.ZeroInit()
|
26
|
-
for size in [(100,), (10, 20), (10, 20, 30)]:
|
27
|
-
weights = init(size)
|
28
|
-
assert weights.shape == size
|
29
|
-
|
30
|
-
|
31
|
-
class TestOneInit(unittest.TestCase):
|
32
|
-
def test_one_init(self):
|
33
|
-
for size in [(100,), (10, 20), (10, 20, 30)]:
|
34
|
-
for value in [0., 1., -1.]:
|
35
|
-
init = brainstate.init.Constant(value=value)
|
36
|
-
weights = init(size)
|
37
|
-
assert weights.shape == size
|
38
|
-
assert (weights == value).all()
|
39
|
-
|
40
|
-
|
41
|
-
class TestIdentityInit(unittest.TestCase):
|
42
|
-
def test_identity_init(self):
|
43
|
-
for size in [(100,), (10, 20)]:
|
44
|
-
for value in [0., 1., -1.]:
|
45
|
-
init = brainstate.init.Identity(value=value)
|
46
|
-
weights = init(size)
|
47
|
-
if len(size) == 1:
|
48
|
-
assert weights.shape == (size[0], size[0])
|
49
|
-
else:
|
50
|
-
assert weights.shape == size
|