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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (112) hide show
  1. brainstate/__init__.py +169 -169
  2. brainstate/_compatible_import.py +340 -340
  3. brainstate/_compatible_import_test.py +681 -681
  4. brainstate/_deprecation.py +210 -210
  5. brainstate/_deprecation_test.py +2319 -2319
  6. brainstate/_error.py +45 -45
  7. brainstate/_state.py +1652 -1652
  8. brainstate/_state_test.py +52 -52
  9. brainstate/_utils.py +47 -47
  10. brainstate/environ.py +1495 -1495
  11. brainstate/environ_test.py +1223 -1223
  12. brainstate/graph/__init__.py +22 -22
  13. brainstate/graph/_node.py +240 -240
  14. brainstate/graph/_node_test.py +589 -589
  15. brainstate/graph/_operation.py +1624 -1624
  16. brainstate/graph/_operation_test.py +1147 -1147
  17. brainstate/mixin.py +1433 -1433
  18. brainstate/mixin_test.py +1017 -1017
  19. brainstate/nn/__init__.py +137 -137
  20. brainstate/nn/_activations.py +1100 -1100
  21. brainstate/nn/_activations_test.py +354 -354
  22. brainstate/nn/_collective_ops.py +633 -633
  23. brainstate/nn/_collective_ops_test.py +774 -774
  24. brainstate/nn/_common.py +226 -226
  25. brainstate/nn/_common_test.py +154 -154
  26. brainstate/nn/_conv.py +2010 -2010
  27. brainstate/nn/_conv_test.py +849 -849
  28. brainstate/nn/_delay.py +575 -575
  29. brainstate/nn/_delay_test.py +243 -243
  30. brainstate/nn/_dropout.py +618 -618
  31. brainstate/nn/_dropout_test.py +477 -477
  32. brainstate/nn/_dynamics.py +1267 -1267
  33. brainstate/nn/_dynamics_test.py +67 -67
  34. brainstate/nn/_elementwise.py +1298 -1298
  35. brainstate/nn/_elementwise_test.py +829 -829
  36. brainstate/nn/_embedding.py +408 -408
  37. brainstate/nn/_embedding_test.py +156 -156
  38. brainstate/nn/_event_fixedprob.py +233 -233
  39. brainstate/nn/_event_fixedprob_test.py +115 -115
  40. brainstate/nn/_event_linear.py +83 -83
  41. brainstate/nn/_event_linear_test.py +121 -121
  42. brainstate/nn/_exp_euler.py +254 -254
  43. brainstate/nn/_exp_euler_test.py +377 -377
  44. brainstate/nn/_linear.py +744 -744
  45. brainstate/nn/_linear_test.py +475 -475
  46. brainstate/nn/_metrics.py +1070 -1070
  47. brainstate/nn/_metrics_test.py +611 -611
  48. brainstate/nn/_module.py +384 -384
  49. brainstate/nn/_module_test.py +40 -40
  50. brainstate/nn/_normalizations.py +1334 -1334
  51. brainstate/nn/_normalizations_test.py +699 -699
  52. brainstate/nn/_paddings.py +1020 -1020
  53. brainstate/nn/_paddings_test.py +722 -722
  54. brainstate/nn/_poolings.py +2239 -2239
  55. brainstate/nn/_poolings_test.py +952 -952
  56. brainstate/nn/_rnns.py +946 -946
  57. brainstate/nn/_rnns_test.py +592 -592
  58. brainstate/nn/_utils.py +216 -216
  59. brainstate/nn/_utils_test.py +401 -401
  60. brainstate/nn/init.py +809 -809
  61. brainstate/nn/init_test.py +180 -180
  62. brainstate/random/__init__.py +270 -270
  63. brainstate/random/_rand_funs.py +3938 -3938
  64. brainstate/random/_rand_funs_test.py +640 -640
  65. brainstate/random/_rand_seed.py +675 -675
  66. brainstate/random/_rand_seed_test.py +48 -48
  67. brainstate/random/_rand_state.py +1617 -1617
  68. brainstate/random/_rand_state_test.py +551 -551
  69. brainstate/transform/__init__.py +59 -59
  70. brainstate/transform/_ad_checkpoint.py +176 -176
  71. brainstate/transform/_ad_checkpoint_test.py +49 -49
  72. brainstate/transform/_autograd.py +1025 -1025
  73. brainstate/transform/_autograd_test.py +1289 -1289
  74. brainstate/transform/_conditions.py +316 -316
  75. brainstate/transform/_conditions_test.py +220 -220
  76. brainstate/transform/_error_if.py +94 -94
  77. brainstate/transform/_error_if_test.py +52 -52
  78. brainstate/transform/_eval_shape.py +145 -145
  79. brainstate/transform/_eval_shape_test.py +38 -38
  80. brainstate/transform/_jit.py +399 -399
  81. brainstate/transform/_jit_test.py +143 -143
  82. brainstate/transform/_loop_collect_return.py +675 -675
  83. brainstate/transform/_loop_collect_return_test.py +58 -58
  84. brainstate/transform/_loop_no_collection.py +283 -283
  85. brainstate/transform/_loop_no_collection_test.py +50 -50
  86. brainstate/transform/_make_jaxpr.py +2016 -2016
  87. brainstate/transform/_make_jaxpr_test.py +1510 -1510
  88. brainstate/transform/_mapping.py +529 -529
  89. brainstate/transform/_mapping_test.py +194 -194
  90. brainstate/transform/_progress_bar.py +255 -255
  91. brainstate/transform/_random.py +171 -171
  92. brainstate/transform/_unvmap.py +256 -256
  93. brainstate/transform/_util.py +286 -286
  94. brainstate/typing.py +837 -837
  95. brainstate/typing_test.py +780 -780
  96. brainstate/util/__init__.py +27 -27
  97. brainstate/util/_others.py +1024 -1024
  98. brainstate/util/_others_test.py +962 -962
  99. brainstate/util/_pretty_pytree.py +1301 -1301
  100. brainstate/util/_pretty_pytree_test.py +675 -675
  101. brainstate/util/_pretty_repr.py +462 -462
  102. brainstate/util/_pretty_repr_test.py +696 -696
  103. brainstate/util/filter.py +945 -945
  104. brainstate/util/filter_test.py +911 -911
  105. brainstate/util/struct.py +910 -910
  106. brainstate/util/struct_test.py +602 -602
  107. {brainstate-0.2.0.dist-info → brainstate-0.2.1.dist-info}/METADATA +108 -108
  108. brainstate-0.2.1.dist-info/RECORD +111 -0
  109. {brainstate-0.2.0.dist-info → brainstate-0.2.1.dist-info}/licenses/LICENSE +202 -202
  110. brainstate-0.2.0.dist-info/RECORD +0 -111
  111. {brainstate-0.2.0.dist-info → brainstate-0.2.1.dist-info}/WHEEL +0 -0
  112. {brainstate-0.2.0.dist-info → brainstate-0.2.1.dist-info}/top_level.txt +0 -0
@@ -1,171 +1,171 @@
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
- import functools
17
- from typing import Callable, Sequence, Union
18
-
19
- from brainstate._utils import set_module_as
20
- from brainstate.random import DEFAULT, RandomState
21
- from brainstate.typing import Missing
22
- from brainstate.util import PrettyObject
23
-
24
- __all__ = [
25
- 'restore_rngs'
26
- ]
27
-
28
-
29
- class RngRestore(PrettyObject):
30
- """
31
- Manage backing up and restoring multiple random states.
32
-
33
- Parameters
34
- ----------
35
- rngs : Sequence[RandomState]
36
- Sequence of :class:`~brainstate.random.RandomState` instances whose
37
- states should be captured and restored.
38
-
39
- Attributes
40
- ----------
41
- rngs : Sequence[RandomState]
42
- Managed random-state instances.
43
- rng_keys : list
44
- Cached keys captured by :meth:`backup` until :meth:`restore` runs.
45
-
46
- Examples
47
- --------
48
- .. code-block:: python
49
-
50
- >>> import brainstate
51
- >>>
52
- >>> rng = brainstate.random.RandomState(0)
53
- >>> restorer = brainstate.transform.RngRestore([rng])
54
- >>> restorer.backup()
55
- >>> _ = rng.random()
56
- >>> restorer.restore()
57
- """
58
- __module__ = 'brainstate.transform'
59
-
60
- def __init__(self, rngs: Sequence[RandomState]):
61
- """
62
- Initialize a restorer for the provided random states.
63
-
64
- Parameters
65
- ----------
66
- rngs : Sequence[RandomState]
67
- Random states that will be backed up and restored.
68
- """
69
- self.rngs: Sequence[RandomState] = rngs
70
- self.rng_keys = []
71
-
72
- def backup(self):
73
- """
74
- Cache the current key for each managed random state.
75
-
76
- Notes
77
- -----
78
- The cached keys persist until :meth:`restore` is called, after which the
79
- internal cache is cleared.
80
- """
81
- self.rng_keys = [rng.value for rng in self.rngs]
82
-
83
- def restore(self):
84
- """
85
- Restore each random state to the cached key.
86
-
87
- Raises
88
- ------
89
- ValueError
90
- Raised when the number of stored keys does not match ``rngs``.
91
- """
92
- if len(self.rng_keys) != len(self.rngs):
93
- raise ValueError('The number of random keys does not match the number of random states.')
94
- for rng, key in zip(self.rngs, self.rng_keys):
95
- rng.restore_value(key)
96
- self.rng_keys.clear()
97
-
98
-
99
- def _rng_backup(
100
- fn: Callable,
101
- rngs: Union[RandomState, Sequence[RandomState]]
102
- ) -> Callable:
103
- rng_restorer = RngRestore(rngs)
104
-
105
- @functools.wraps(fn)
106
- def wrapper(*args, **kwargs):
107
- # backup the random state
108
- rng_restorer.backup()
109
- # call the function
110
- out = fn(*args, **kwargs)
111
- # restore the random state
112
- rng_restorer.restore()
113
- return out
114
-
115
- return wrapper
116
-
117
-
118
- @set_module_as('brainstate.transform')
119
- def restore_rngs(
120
- fn: Callable = Missing(),
121
- rngs: Union[RandomState, Sequence[RandomState]] = DEFAULT,
122
- ) -> Callable:
123
- """
124
- Decorate a function so specified random states are restored after execution.
125
-
126
- Parameters
127
- ----------
128
- fn : Callable, optional
129
- Function to wrap. When omitted, :func:`restore_rngs` returns a decorator
130
- preconfigured with ``rngs``.
131
- rngs : Union[RandomState, Sequence[RandomState]], optional
132
- Random states whose keys should be backed up before running ``fn`` and
133
- restored afterwards. Defaults to :data:`brainstate.random.DEFAULT`.
134
-
135
- Returns
136
- -------
137
- Callable
138
- Wrapped callable that restores the random state or a partially applied
139
- decorator depending on how :func:`restore_rngs` is used.
140
-
141
- Raises
142
- ------
143
- AssertionError
144
- If ``rngs`` is neither a :class:`~brainstate.random.RandomState` instance nor
145
- a sequence of such instances.
146
-
147
- Examples
148
- --------
149
- .. code-block:: python
150
-
151
- >>> import brainstate
152
- >>>
153
- >>> rng = brainstate.random.RandomState(0)
154
- >>>
155
- >>> @brainstate.transform.restore_rngs(rngs=rng)
156
- ... def sample_pair():
157
- ... first = rng.random()
158
- ... second = rng.random()
159
- ... return first, second
160
- >>>
161
- >>> assert sample_pair()[0] == sample_pair()[0]
162
- """
163
- if isinstance(fn, Missing):
164
- return functools.partial(restore_rngs, rngs=rngs)
165
-
166
- if isinstance(rngs, RandomState):
167
- rngs = [rngs]
168
- assert isinstance(rngs, Sequence), 'rngs must be a RandomState or a sequence of RandomState instances.'
169
- for rng in rngs:
170
- assert isinstance(rng, RandomState), 'rngs must be a RandomState or a sequence of RandomState instances.'
171
- return _rng_backup(fn, rngs=rngs)
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
+ import functools
17
+ from typing import Callable, Sequence, Union
18
+
19
+ from brainstate._utils import set_module_as
20
+ from brainstate.random import DEFAULT, RandomState
21
+ from brainstate.typing import Missing
22
+ from brainstate.util import PrettyObject
23
+
24
+ __all__ = [
25
+ 'restore_rngs'
26
+ ]
27
+
28
+
29
+ class RngRestore(PrettyObject):
30
+ """
31
+ Manage backing up and restoring multiple random states.
32
+
33
+ Parameters
34
+ ----------
35
+ rngs : Sequence[RandomState]
36
+ Sequence of :class:`~brainstate.random.RandomState` instances whose
37
+ states should be captured and restored.
38
+
39
+ Attributes
40
+ ----------
41
+ rngs : Sequence[RandomState]
42
+ Managed random-state instances.
43
+ rng_keys : list
44
+ Cached keys captured by :meth:`backup` until :meth:`restore` runs.
45
+
46
+ Examples
47
+ --------
48
+ .. code-block:: python
49
+
50
+ >>> import brainstate
51
+ >>>
52
+ >>> rng = brainstate.random.RandomState(0)
53
+ >>> restorer = brainstate.transform.RngRestore([rng])
54
+ >>> restorer.backup()
55
+ >>> _ = rng.random()
56
+ >>> restorer.restore()
57
+ """
58
+ __module__ = 'brainstate.transform'
59
+
60
+ def __init__(self, rngs: Sequence[RandomState]):
61
+ """
62
+ Initialize a restorer for the provided random states.
63
+
64
+ Parameters
65
+ ----------
66
+ rngs : Sequence[RandomState]
67
+ Random states that will be backed up and restored.
68
+ """
69
+ self.rngs: Sequence[RandomState] = rngs
70
+ self.rng_keys = []
71
+
72
+ def backup(self):
73
+ """
74
+ Cache the current key for each managed random state.
75
+
76
+ Notes
77
+ -----
78
+ The cached keys persist until :meth:`restore` is called, after which the
79
+ internal cache is cleared.
80
+ """
81
+ self.rng_keys = [rng.value for rng in self.rngs]
82
+
83
+ def restore(self):
84
+ """
85
+ Restore each random state to the cached key.
86
+
87
+ Raises
88
+ ------
89
+ ValueError
90
+ Raised when the number of stored keys does not match ``rngs``.
91
+ """
92
+ if len(self.rng_keys) != len(self.rngs):
93
+ raise ValueError('The number of random keys does not match the number of random states.')
94
+ for rng, key in zip(self.rngs, self.rng_keys):
95
+ rng.restore_value(key)
96
+ self.rng_keys.clear()
97
+
98
+
99
+ def _rng_backup(
100
+ fn: Callable,
101
+ rngs: Union[RandomState, Sequence[RandomState]]
102
+ ) -> Callable:
103
+ rng_restorer = RngRestore(rngs)
104
+
105
+ @functools.wraps(fn)
106
+ def wrapper(*args, **kwargs):
107
+ # backup the random state
108
+ rng_restorer.backup()
109
+ # call the function
110
+ out = fn(*args, **kwargs)
111
+ # restore the random state
112
+ rng_restorer.restore()
113
+ return out
114
+
115
+ return wrapper
116
+
117
+
118
+ @set_module_as('brainstate.transform')
119
+ def restore_rngs(
120
+ fn: Callable = Missing(),
121
+ rngs: Union[RandomState, Sequence[RandomState]] = DEFAULT,
122
+ ) -> Callable:
123
+ """
124
+ Decorate a function so specified random states are restored after execution.
125
+
126
+ Parameters
127
+ ----------
128
+ fn : Callable, optional
129
+ Function to wrap. When omitted, :func:`restore_rngs` returns a decorator
130
+ preconfigured with ``rngs``.
131
+ rngs : Union[RandomState, Sequence[RandomState]], optional
132
+ Random states whose keys should be backed up before running ``fn`` and
133
+ restored afterwards. Defaults to :data:`brainstate.random.DEFAULT`.
134
+
135
+ Returns
136
+ -------
137
+ Callable
138
+ Wrapped callable that restores the random state or a partially applied
139
+ decorator depending on how :func:`restore_rngs` is used.
140
+
141
+ Raises
142
+ ------
143
+ AssertionError
144
+ If ``rngs`` is neither a :class:`~brainstate.random.RandomState` instance nor
145
+ a sequence of such instances.
146
+
147
+ Examples
148
+ --------
149
+ .. code-block:: python
150
+
151
+ >>> import brainstate
152
+ >>>
153
+ >>> rng = brainstate.random.RandomState(0)
154
+ >>>
155
+ >>> @brainstate.transform.restore_rngs(rngs=rng)
156
+ ... def sample_pair():
157
+ ... first = rng.random()
158
+ ... second = rng.random()
159
+ ... return first, second
160
+ >>>
161
+ >>> assert sample_pair()[0] == sample_pair()[0]
162
+ """
163
+ if isinstance(fn, Missing):
164
+ return functools.partial(restore_rngs, rngs=rngs)
165
+
166
+ if isinstance(rngs, RandomState):
167
+ rngs = [rngs]
168
+ assert isinstance(rngs, Sequence), 'rngs must be a RandomState or a sequence of RandomState instances.'
169
+ for rng in rngs:
170
+ assert isinstance(rng, RandomState), 'rngs must be a RandomState or a sequence of RandomState instances.'
171
+ return _rng_backup(fn, rngs=rngs)