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,29 +1,22 @@
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
-
17
- from ._graph_node import Node, Dict, List, Sequential
18
- from ._graph_operation import (
19
- pop_states, nodes, states, treefy_states, update_states, flatten, unflatten,
20
- treefy_split, treefy_merge, iter_leaf, iter_node, clone, graphdef,
21
- call, RefMap, GraphDef, NodeRef, NodeDef
22
- )
23
-
24
- __all__ = [
25
- 'Node', 'Dict', 'List', 'Sequential',
26
- 'pop_states', 'nodes', 'states', 'treefy_states', 'update_states', 'flatten', 'unflatten',
27
- 'treefy_split', 'treefy_merge', 'iter_leaf', 'iter_node', 'clone', 'graphdef',
28
- 'call', 'RefMap', 'GraphDef', 'NodeRef', 'NodeDef',
29
- ]
1
+ # Copyright 2024 BrainX Ecosystem Limited. All Rights Reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+ # ==============================================================================
15
+
16
+
17
+ from ._node import Node
18
+ from ._operation import *
19
+ from ._operation import __all__ as operation_all
20
+
21
+ __all__ = ['Node'] + operation_all
22
+ del operation_all
@@ -0,0 +1,240 @@
1
+ # The file is adapted from the Flax library (https://github.com/google/flax).
2
+ # The credit should go to the Flax authors.
3
+ #
4
+ # Copyright 2024 The Flax Authors.
5
+ #
6
+ # Licensed under the Apache License, Version 2.0 (the "License");
7
+ # you may not use this file except in compliance with the License.
8
+ # You may obtain a copy of the License at
9
+ #
10
+ # http://www.apache.org/licenses/LICENSE-2.0
11
+ #
12
+ # Unless required by applicable law or agreed to in writing, software
13
+ # distributed under the License is distributed on an "AS IS" BASIS,
14
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15
+ # See the License for the specific language governing permissions and
16
+ # limitations under the License.
17
+
18
+ from abc import ABCMeta
19
+ from copy import deepcopy
20
+ from typing import Any, Type, TypeVar, Tuple, TYPE_CHECKING
21
+
22
+ from brainstate._state import State, TreefyState
23
+ from brainstate.typing import Key
24
+ from brainstate.util._pretty_pytree import PrettyObject
25
+ from ._operation import register_graph_node_type, treefy_split, treefy_merge
26
+
27
+ __all__ = [
28
+ 'Node',
29
+ ]
30
+
31
+ G = TypeVar('G', bound='Node')
32
+ A = TypeVar('A')
33
+
34
+
35
+ class GraphNodeMeta(ABCMeta):
36
+ if not TYPE_CHECKING:
37
+ def __call__(cls, *args, **kwargs) -> Any:
38
+ node = cls.__new__(cls, *args, **kwargs)
39
+ node.__init__(*args, **kwargs)
40
+ return node
41
+
42
+
43
+ class Node(PrettyObject, metaclass=GraphNodeMeta):
44
+ """
45
+ Base class for all graph nodes in the BrainState framework.
46
+
47
+ This class serves as the foundation for creating computational graph nodes
48
+ that can be used in neural network architectures and other graph-based
49
+ computations.
50
+
51
+ Attributes
52
+ ----------
53
+ graph_invisible_attrs : tuple
54
+ Tuple of attribute names that should be excluded from graph
55
+ serialization and flattening operations.
56
+
57
+ Methods
58
+ -------
59
+ __deepcopy__(memo=None)
60
+ Creates a deep copy of the node preserving its graph structure
61
+ and state.
62
+
63
+ Notes
64
+ -----
65
+ The class provides the following features:
66
+
67
+ - Automatic registration with the graph system via metaclass
68
+ - Deep copy support for creating independent node instances
69
+ - Pretty printing for better debugging and visualization
70
+ - State management integration with TreefyState
71
+ - Attribute visibility control via graph_invisible_attrs
72
+
73
+ Examples
74
+ --------
75
+ .. code-block:: python
76
+
77
+ >>> from copy import deepcopy
78
+ >>> class MyNode(Node):
79
+ ... def __init__(self, value):
80
+ ... self.value = value
81
+ >>> node = MyNode(10)
82
+ >>> copied_node = deepcopy(node)
83
+ >>> print(node.value)
84
+ 10
85
+ """
86
+ __module__ = 'brainstate.graph'
87
+
88
+ graph_invisible_attrs = ()
89
+
90
+ def __init_subclass__(cls) -> None:
91
+ super().__init_subclass__()
92
+
93
+ register_graph_node_type(
94
+ type=cls,
95
+ flatten=_node_flatten,
96
+ set_key=_node_set_key,
97
+ pop_key=_node_pop_key,
98
+ create_empty=_node_create_empty,
99
+ clear=_node_clear,
100
+ )
101
+
102
+ def __deepcopy__(self: G, memo=None) -> G:
103
+ graphdef, state = treefy_split(self)
104
+ graphdef = deepcopy(graphdef)
105
+ state = deepcopy(state)
106
+ return treefy_merge(graphdef, state)
107
+
108
+
109
+ # -------------------------------
110
+ # Graph Definition
111
+ # -------------------------------
112
+
113
+
114
+ def _node_flatten(node: Node) -> Tuple[Tuple[Tuple[str, Any], ...], Tuple[Type]]:
115
+ """
116
+ Flatten a node into its constituent parts for serialization.
117
+
118
+ Parameters
119
+ ----------
120
+ node : Node
121
+ The Node instance to flatten.
122
+
123
+ Returns
124
+ -------
125
+ tuple
126
+ A tuple containing:
127
+ - Sorted list of (key, value) pairs for visible attributes
128
+ - Tuple containing the node's type
129
+ """
130
+ graph_invisible_attrs = getattr(node, 'graph_invisible_attrs', ())
131
+ # graph_invisible_attrs = tuple(graph_invisible_attrs) + ('_trace_state',)
132
+ nodes = sorted(
133
+ (key, value) for key, value in vars(node).items()
134
+ if (key not in graph_invisible_attrs)
135
+ )
136
+ return nodes, (type(node),)
137
+
138
+
139
+ def _node_set_key(node: Node, key: Key, value: Any) -> None:
140
+ """
141
+ Set an attribute on a node with special handling for State objects.
142
+
143
+ Parameters
144
+ ----------
145
+ node : Node
146
+ The Node instance to modify.
147
+ key : Key
148
+ The attribute name to set.
149
+ value : Any
150
+ The value to set.
151
+
152
+ Raises
153
+ ------
154
+ KeyError
155
+ If the key is not a string.
156
+
157
+ Notes
158
+ -----
159
+ If the attribute already exists as a State object and the new value
160
+ is a TreefyState, the state is updated via reference rather than
161
+ replaced.
162
+ """
163
+ if not isinstance(key, str):
164
+ raise KeyError(f'Invalid key: {key!r}')
165
+ elif (
166
+ hasattr(node, key)
167
+ and isinstance(state := getattr(node, key), State)
168
+ and isinstance(value, TreefyState)
169
+ ):
170
+ state.update_from_ref(value)
171
+ else:
172
+ setattr(node, key, value)
173
+
174
+
175
+ def _node_pop_key(node: Node, key: Key) -> Any:
176
+ """
177
+ Remove and return an attribute from a node.
178
+
179
+ Parameters
180
+ ----------
181
+ node : Node
182
+ The Node instance to modify.
183
+ key : Key
184
+ The attribute name to remove.
185
+
186
+ Returns
187
+ -------
188
+ Any
189
+ The value of the removed attribute.
190
+
191
+ Raises
192
+ ------
193
+ KeyError
194
+ If the key is not a string.
195
+ """
196
+ if not isinstance(key, str):
197
+ raise KeyError(f'Invalid key: {key!r}')
198
+ return vars(node).pop(key)
199
+
200
+
201
+ def _node_create_empty(static: tuple[Type[G], ...]) -> G:
202
+ """
203
+ Create an empty node instance without calling __init__.
204
+
205
+ Parameters
206
+ ----------
207
+ static : tuple[Type[G], ...]
208
+ Tuple containing the node type to instantiate.
209
+
210
+ Returns
211
+ -------
212
+ G
213
+ A new uninitialized instance of the node type.
214
+
215
+ Notes
216
+ -----
217
+ This function is used internally by the graph system to create
218
+ nodes without invoking their initialization logic.
219
+ """
220
+ node_type, = static
221
+ node = object.__new__(node_type)
222
+ return node
223
+
224
+
225
+ def _node_clear(node: Node) -> None:
226
+ """
227
+ Clear all attributes from a node.
228
+
229
+ Parameters
230
+ ----------
231
+ node : Node
232
+ The Node instance to clear.
233
+
234
+ Notes
235
+ -----
236
+ This removes all attributes from the node's instance dictionary,
237
+ effectively resetting it to an empty state.
238
+ """
239
+ module_vars = vars(node)
240
+ module_vars.clear()