brainstate 0.1.10__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.
Files changed (163) hide show
  1. brainstate/__init__.py +130 -19
  2. brainstate/_compatible_import.py +201 -9
  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} +10 -20
  7. brainstate/_state.py +94 -47
  8. brainstate/_state_test.py +1 -1
  9. brainstate/_utils.py +1 -1
  10. brainstate/environ.py +1279 -347
  11. brainstate/environ_test.py +1187 -26
  12. brainstate/graph/__init__.py +6 -13
  13. brainstate/graph/_node.py +240 -0
  14. brainstate/graph/_node_test.py +589 -0
  15. brainstate/graph/{_graph_operation.py → _operation.py} +632 -746
  16. brainstate/graph/_operation_test.py +1147 -0
  17. brainstate/mixin.py +1209 -141
  18. brainstate/mixin_test.py +991 -51
  19. brainstate/nn/__init__.py +74 -72
  20. brainstate/nn/_activations.py +587 -295
  21. brainstate/nn/_activations_test.py +109 -86
  22. brainstate/nn/_collective_ops.py +393 -274
  23. brainstate/nn/_collective_ops_test.py +746 -15
  24. brainstate/nn/_common.py +114 -66
  25. brainstate/nn/_common_test.py +154 -0
  26. brainstate/nn/_conv.py +1652 -143
  27. brainstate/nn/_conv_test.py +838 -227
  28. brainstate/nn/_delay.py +15 -28
  29. brainstate/nn/_delay_test.py +25 -20
  30. brainstate/nn/_dropout.py +359 -167
  31. brainstate/nn/_dropout_test.py +429 -52
  32. brainstate/nn/_dynamics.py +14 -90
  33. brainstate/nn/_dynamics_test.py +1 -12
  34. brainstate/nn/_elementwise.py +492 -313
  35. brainstate/nn/_elementwise_test.py +806 -145
  36. brainstate/nn/_embedding.py +369 -19
  37. brainstate/nn/_embedding_test.py +156 -0
  38. brainstate/nn/{_fixedprob.py → _event_fixedprob.py} +10 -16
  39. brainstate/nn/{_fixedprob_test.py → _event_fixedprob_test.py} +6 -5
  40. brainstate/nn/{_linear_mv.py → _event_linear.py} +2 -2
  41. brainstate/nn/{_linear_mv_test.py → _event_linear_test.py} +6 -5
  42. brainstate/nn/_exp_euler.py +200 -38
  43. brainstate/nn/_exp_euler_test.py +350 -8
  44. brainstate/nn/_linear.py +391 -71
  45. brainstate/nn/_linear_test.py +427 -59
  46. brainstate/nn/_metrics.py +1070 -0
  47. brainstate/nn/_metrics_test.py +611 -0
  48. brainstate/nn/_module.py +10 -3
  49. brainstate/nn/_module_test.py +1 -1
  50. brainstate/nn/_normalizations.py +688 -329
  51. brainstate/nn/_normalizations_test.py +663 -37
  52. brainstate/nn/_paddings.py +1020 -0
  53. brainstate/nn/_paddings_test.py +723 -0
  54. brainstate/nn/_poolings.py +1404 -342
  55. brainstate/nn/_poolings_test.py +828 -92
  56. brainstate/nn/{_rate_rnns.py → _rnns.py} +446 -54
  57. brainstate/nn/_rnns_test.py +593 -0
  58. brainstate/nn/_utils.py +132 -5
  59. brainstate/nn/_utils_test.py +402 -0
  60. brainstate/{init/_random_inits.py → nn/init.py} +301 -45
  61. brainstate/{init/_random_inits_test.py → nn/init_test.py} +51 -20
  62. brainstate/random/__init__.py +247 -1
  63. brainstate/random/_rand_funs.py +668 -346
  64. brainstate/random/_rand_funs_test.py +74 -1
  65. brainstate/random/_rand_seed.py +541 -76
  66. brainstate/random/_rand_seed_test.py +1 -1
  67. brainstate/random/_rand_state.py +601 -393
  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 +1 -1
  72. brainstate/{augment → transform}/_autograd.py +360 -113
  73. brainstate/{augment → transform}/_autograd_test.py +2 -2
  74. brainstate/transform/_conditions.py +316 -0
  75. brainstate/{compile → transform}/_conditions_test.py +11 -11
  76. brainstate/{compile → transform}/_error_if.py +22 -20
  77. brainstate/{compile → transform}/_error_if_test.py +1 -1
  78. brainstate/transform/_eval_shape.py +145 -0
  79. brainstate/{augment → transform}/_eval_shape_test.py +1 -1
  80. brainstate/{compile → transform}/_jit.py +99 -46
  81. brainstate/{compile → transform}/_jit_test.py +3 -3
  82. brainstate/{compile → transform}/_loop_collect_return.py +219 -80
  83. brainstate/{compile → transform}/_loop_collect_return_test.py +1 -1
  84. brainstate/{compile → transform}/_loop_no_collection.py +133 -34
  85. brainstate/{compile → transform}/_loop_no_collection_test.py +2 -2
  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 +78 -25
  91. brainstate/{augment → transform}/_random.py +65 -45
  92. brainstate/{compile → transform}/_unvmap.py +102 -5
  93. brainstate/transform/_util.py +286 -0
  94. brainstate/typing.py +594 -61
  95. brainstate/typing_test.py +780 -0
  96. brainstate/util/__init__.py +9 -32
  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} +161 -27
  102. brainstate/util/_pretty_repr_test.py +696 -0
  103. brainstate/util/filter.py +557 -81
  104. brainstate/util/filter_test.py +912 -0
  105. brainstate/util/struct.py +769 -382
  106. brainstate/util/struct_test.py +602 -0
  107. {brainstate-0.1.10.dist-info → brainstate-0.2.0.dist-info}/METADATA +34 -17
  108. brainstate-0.2.0.dist-info/RECORD +111 -0
  109. brainstate/augment/__init__.py +0 -30
  110. brainstate/augment/_eval_shape.py +0 -99
  111. brainstate/augment/_mapping.py +0 -1060
  112. brainstate/augment/_mapping_test.py +0 -597
  113. brainstate/compile/__init__.py +0 -38
  114. brainstate/compile/_ad_checkpoint.py +0 -204
  115. brainstate/compile/_conditions.py +0 -256
  116. brainstate/compile/_make_jaxpr.py +0 -888
  117. brainstate/compile/_make_jaxpr_test.py +0 -156
  118. brainstate/compile/_util.py +0 -147
  119. brainstate/functional/__init__.py +0 -27
  120. brainstate/graph/_graph_node.py +0 -244
  121. brainstate/graph/_graph_node_test.py +0 -73
  122. brainstate/graph/_graph_operation_test.py +0 -563
  123. brainstate/init/__init__.py +0 -26
  124. brainstate/init/_base.py +0 -52
  125. brainstate/init/_generic.py +0 -244
  126. brainstate/init/_regular_inits.py +0 -105
  127. brainstate/init/_regular_inits_test.py +0 -50
  128. brainstate/nn/_inputs.py +0 -608
  129. brainstate/nn/_ltp.py +0 -28
  130. brainstate/nn/_neuron.py +0 -705
  131. brainstate/nn/_neuron_test.py +0 -161
  132. brainstate/nn/_others.py +0 -46
  133. brainstate/nn/_projection.py +0 -486
  134. brainstate/nn/_rate_rnns_test.py +0 -63
  135. brainstate/nn/_readout.py +0 -209
  136. brainstate/nn/_readout_test.py +0 -53
  137. brainstate/nn/_stp.py +0 -236
  138. brainstate/nn/_synapse.py +0 -505
  139. brainstate/nn/_synapse_test.py +0 -131
  140. brainstate/nn/_synaptic_projection.py +0 -423
  141. brainstate/nn/_synouts.py +0 -162
  142. brainstate/nn/_synouts_test.py +0 -57
  143. brainstate/nn/metrics.py +0 -388
  144. brainstate/optim/__init__.py +0 -38
  145. brainstate/optim/_base.py +0 -64
  146. brainstate/optim/_lr_scheduler.py +0 -448
  147. brainstate/optim/_lr_scheduler_test.py +0 -50
  148. brainstate/optim/_optax_optimizer.py +0 -152
  149. brainstate/optim/_optax_optimizer_test.py +0 -53
  150. brainstate/optim/_sgd_optimizer.py +0 -1104
  151. brainstate/random/_random_for_unit.py +0 -52
  152. brainstate/surrogate.py +0 -1957
  153. brainstate/transform.py +0 -23
  154. brainstate/util/caller.py +0 -98
  155. brainstate/util/others.py +0 -540
  156. brainstate/util/pretty_pytree.py +0 -945
  157. brainstate/util/pretty_pytree_test.py +0 -159
  158. brainstate/util/pretty_table.py +0 -2954
  159. brainstate/util/scaling.py +0 -258
  160. brainstate-0.1.10.dist-info/RECORD +0 -130
  161. {brainstate-0.1.10.dist-info → brainstate-0.2.0.dist-info}/WHEEL +0 -0
  162. {brainstate-0.1.10.dist-info → brainstate-0.2.0.dist-info}/licenses/LICENSE +0 -0
  163. {brainstate-0.1.10.dist-info → brainstate-0.2.0.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,210 @@
1
+ # Copyright 2025 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
+ Deprecation utilities for brainstate.
18
+ """
19
+
20
+ import warnings
21
+
22
+
23
+ class DeprecatedModule:
24
+ """
25
+ A proxy class that mimics a module but shows deprecation warnings when accessed.
26
+
27
+ This class allows for smooth deprecation of modules by forwarding attribute
28
+ access to a replacement module while showing appropriate warnings.
29
+ """
30
+
31
+ def __init__(
32
+ self,
33
+ deprecated_name,
34
+ replacement_module,
35
+ replacement_name,
36
+ version="0.1.11",
37
+ removal_version=None,
38
+ scoped_apis=None
39
+ ):
40
+ """
41
+ Initialize a deprecated module proxy.
42
+
43
+ Args:
44
+ deprecated_name: Name of the deprecated module (e.g., 'brainstate.augment')
45
+ replacement_module: The module to forward calls to
46
+ replacement_name: Name of the replacement module (e.g., 'brainstate.transform')
47
+ version: Version when deprecation started
48
+ removal_version: Version when module will be removed (optional)
49
+ scoped_apis: Dict mapping API names to their locations, or list of API names
50
+ """
51
+ self._deprecated_name = deprecated_name
52
+ self._replacement_module = replacement_module
53
+ self._replacement_name = replacement_name
54
+ self._version = version
55
+ self._removal_version = removal_version
56
+ self._warned_attrs = {}
57
+ self._scoped_apis = scoped_apis or {}
58
+
59
+ # Set module-like attributes
60
+ self.__name__ = deprecated_name
61
+ self.__doc__ = f"DEPRECATED: {deprecated_name} is deprecated. Use {replacement_name} instead."
62
+
63
+ # Handle scoped APIs
64
+ if isinstance(scoped_apis, dict):
65
+ # Dict mapping API names to their import locations
66
+ self._api_mapping = scoped_apis
67
+ self.__all__ = list(scoped_apis.keys())
68
+ elif isinstance(scoped_apis, (list, tuple)):
69
+ # List of API names to expose from replacement module
70
+ self._api_mapping = {name: replacement_module for name in scoped_apis}
71
+ self.__all__ = list(scoped_apis)
72
+ else:
73
+ # No scoping - use entire replacement module
74
+ self._api_mapping = {}
75
+ if hasattr(replacement_module, '__all__'):
76
+ self.__all__ = replacement_module.__all__
77
+
78
+ @property
79
+ def replacement_module(self):
80
+ if isinstance(self._replacement_module, str):
81
+ # Lazy import of replacement module
82
+ import importlib
83
+ self._replacement_module = importlib.import_module(self._replacement_module)
84
+ return self._replacement_module
85
+
86
+ def _warn_deprecation(self, attr_name=None):
87
+ """Show deprecation warning for module or attribute access."""
88
+ # Only warn once per attribute to avoid spam
89
+ warn_key = attr_name or '__module__'
90
+ if warn_key in self._warned_attrs:
91
+ return
92
+ self._warned_attrs[warn_key] = True
93
+
94
+ if attr_name:
95
+ message = (
96
+ f"Accessing '{attr_name}' from '{self._deprecated_name}' is deprecated "
97
+ f"and will be removed in a future version. "
98
+ f"Use '{self._replacement_name}.{attr_name}' instead."
99
+ )
100
+ else:
101
+ message = (
102
+ f"The '{self._deprecated_name}' module is deprecated "
103
+ f"and will be removed in a future version. "
104
+ f"Use '{self._replacement_name}' instead."
105
+ )
106
+
107
+ if self._removal_version:
108
+ message += f" It will be removed in version {self._removal_version}."
109
+
110
+ warnings.warn(message, DeprecationWarning, stacklevel=3)
111
+
112
+ def __getattr__(self, name):
113
+ """Forward attribute access to replacement module with deprecation warning."""
114
+ self._warn_deprecation(name)
115
+
116
+ # Check if we have scoped APIs
117
+ if self._api_mapping:
118
+ if name in self._api_mapping:
119
+ # Get from specific location
120
+ source = self._api_mapping[name]
121
+
122
+ if isinstance(source, str):
123
+ # Import from module path
124
+ try:
125
+ # Handle relative imports within brainstate
126
+ if source.startswith('brainstate.'):
127
+ # Import the module dynamically
128
+ import importlib
129
+ module = importlib.import_module(source)
130
+ return getattr(module, name)
131
+ else:
132
+ # For other module paths, use standard import
133
+ module_parts = source.split('.')
134
+ module = __import__(source, fromlist=[name])
135
+ return getattr(module, name)
136
+ except (ImportError, AttributeError) as e:
137
+ # Fallback to replacement module
138
+ try:
139
+ return getattr(self._replacement_module, name)
140
+ except AttributeError:
141
+ raise AttributeError(
142
+ f"Module '{self._deprecated_name}' has no attribute '{name}'. "
143
+ f"Failed to import from '{source}': {e}. "
144
+ f"Check '{self._replacement_name}' for available attributes."
145
+ )
146
+ else:
147
+ # Source is a module object
148
+ try:
149
+ return getattr(source, name)
150
+ except AttributeError:
151
+ # Fallback to replacement module
152
+ return getattr(self._replacement_module, name)
153
+ else:
154
+ # Attribute not in scoped APIs
155
+ available_apis = ', '.join(self.__all__ or [])
156
+ raise AttributeError(
157
+ f"Module '{self._deprecated_name}' has no attribute '{name}'. "
158
+ f"Available attributes: {available_apis}. "
159
+ f"Check '{self._replacement_name}' for more attributes."
160
+ )
161
+
162
+ # Fallback to replacement module for non-scoped access
163
+ try:
164
+ return getattr(self.replacement_module, name)
165
+ except AttributeError:
166
+ raise AttributeError(
167
+ f"Module '{self._deprecated_name}' has no attribute '{name}'. "
168
+ f"Check '{self._replacement_name}' for available attributes."
169
+ )
170
+
171
+ def __dir__(self):
172
+ """Return attributes from replacement module."""
173
+ self._warn_deprecation()
174
+
175
+ if self._api_mapping:
176
+ # Return only scoped APIs plus standard module attributes
177
+ base_attrs = ['__name__', '__doc__', '__all__']
178
+ return base_attrs + list(self._api_mapping.keys())
179
+ else:
180
+ return dir(self.replacement_module)
181
+
182
+ def __repr__(self):
183
+ """Return a deprecation-aware repr."""
184
+ return f"<DeprecatedModule '{self._deprecated_name}' -> '{self._replacement_name}'>"
185
+
186
+
187
+ def create_deprecated_module_proxy(
188
+ deprecated_name,
189
+ replacement_module,
190
+ replacement_name,
191
+ **kwargs
192
+ ):
193
+ """
194
+ Create a deprecated module proxy.
195
+
196
+ Args:
197
+ deprecated_name: Name of the deprecated module
198
+ replacement_module: The module to forward calls to
199
+ replacement_name: Name of the replacement module
200
+ **kwargs: Additional arguments for DeprecatedModule
201
+
202
+ Returns:
203
+ DeprecatedModule proxy instance
204
+ """
205
+ return DeprecatedModule(
206
+ deprecated_name=deprecated_name,
207
+ replacement_module=replacement_module,
208
+ replacement_name=replacement_name,
209
+ **kwargs
210
+ )