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,143 +1,143 @@
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
- from __future__ import annotations
17
-
18
- import unittest
19
-
20
- import jax.numpy as jnp
21
-
22
- import brainstate as bst
23
-
24
-
25
- class TestJIT(unittest.TestCase):
26
- def test_inner_state_are_not_catched(self):
27
- a = bst.State(bst.random.randn(10))
28
-
29
- @bst.compile.jit
30
- def fun1(inp):
31
- a.value += inp
32
-
33
- b = bst.State(bst.random.randn(1))
34
-
35
- def inner_fun(x):
36
- b.value += x
37
-
38
- bst.compile.for_loop(inner_fun, bst.random.randn(100))
39
-
40
- return a.value + b.value
41
-
42
- print(fun1(1.))
43
- key = fun1.stateful_fun.get_arg_cache_key(1.)
44
- self.assertTrue(len(fun1.stateful_fun.get_states_by_cache(key)) == 2)
45
-
46
- x = bst.random.randn(10)
47
- print(fun1(x))
48
- key = fun1.stateful_fun.get_arg_cache_key(x)
49
- self.assertTrue(len(fun1.stateful_fun.get_states_by_cache(key)) == 2)
50
-
51
- def test_kwargs(self):
52
- a = bst.State(bst.random.randn(10))
53
-
54
- @bst.compile.jit
55
- def fun1(inp):
56
- a.value += inp
57
-
58
- b = bst.State(bst.random.randn(1))
59
-
60
- def inner_fun(x):
61
- b.value += x
62
-
63
- bst.compile.for_loop(inner_fun, bst.random.randn(100))
64
-
65
- return a.value + b.value
66
-
67
- # test kwargs
68
- print(fun1(inp=bst.random.randn(10)))
69
-
70
- def test_jit_compile_sensitive_to_input_shape(self):
71
- global_data = [0]
72
-
73
- @bst.compile.jit
74
- def fun1(inp):
75
- global_data[0] += 1
76
- return inp
77
-
78
- print(fun1(1.))
79
- self.assertTrue(global_data[0] == 1)
80
-
81
- print(fun1(2.))
82
- self.assertTrue(global_data[0] == 1)
83
-
84
- print(fun1(bst.random.randn(10)))
85
- self.assertTrue(global_data[0] == 2)
86
-
87
- print(fun1(bst.random.randn(10, 10)))
88
- self.assertTrue(global_data[0] == 3)
89
-
90
- def test_jit_clear_cache(self):
91
- a = bst.State(bst.random.randn(1))
92
- compiling = []
93
-
94
- @bst.compile.jit
95
- def log2(x):
96
- print('compiling')
97
- compiling.append(1)
98
- ln_x = jnp.log(x)
99
- ln_2 = jnp.log(2.0) + a.value
100
- return ln_x / ln_2
101
-
102
- x = bst.random.randn(1)
103
- print(log2(x)) # compiling
104
- self.assertTrue(len(compiling) == 1)
105
- print(log2(x)) # no compiling
106
- self.assertTrue(len(compiling) == 1)
107
-
108
- log2.clear_cache()
109
- print(log2(x)) # compiling
110
- self.assertTrue(len(compiling) == 2)
111
-
112
- def test_jit_attribute_origin_fun(self):
113
- def fun1(x):
114
- return x
115
-
116
- jitted_fun = bst.compile.jit(fun1)
117
- self.assertTrue(jitted_fun.origin_fun is fun1)
118
- self.assertTrue(isinstance(jitted_fun.stateful_fun, bst.compile.StatefulFunction))
119
- self.assertTrue(callable(jitted_fun.jitted_fun))
120
- self.assertTrue(callable(jitted_fun.clear_cache))
121
-
122
- def test_clear_cache(self):
123
- a = bst.State(bst.random.randn(1))
124
-
125
- @bst.compile.jit
126
- def f_jit(x, y):
127
- print('Compiling')
128
- a.value = jnp.sin(x) + jnp.cos(y)
129
-
130
- f_jit(0.5, 1.0)
131
- f_jit.clear_cache()
132
- f_jit(0.5, 1.0)
133
-
134
- def test_cache(self):
135
- @bst.compile.jit
136
- @bst.compile.jit
137
- @bst.compile.jit
138
- def f(a):
139
- print('Compiling')
140
- print(a)
141
- return a + 1
142
-
143
- print(f(1.))
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
+ from __future__ import annotations
17
+
18
+ import unittest
19
+
20
+ import jax.numpy as jnp
21
+
22
+ import brainstate as bst
23
+
24
+
25
+ class TestJIT(unittest.TestCase):
26
+ def test_inner_state_are_not_catched(self):
27
+ a = bst.State(bst.random.randn(10))
28
+
29
+ @bst.compile.jit
30
+ def fun1(inp):
31
+ a.value += inp
32
+
33
+ b = bst.State(bst.random.randn(1))
34
+
35
+ def inner_fun(x):
36
+ b.value += x
37
+
38
+ bst.compile.for_loop(inner_fun, bst.random.randn(100))
39
+
40
+ return a.value + b.value
41
+
42
+ print(fun1(1.))
43
+ key = fun1.stateful_fun.get_arg_cache_key(1.)
44
+ self.assertTrue(len(fun1.stateful_fun.get_states_by_cache(key)) == 2)
45
+
46
+ x = bst.random.randn(10)
47
+ print(fun1(x))
48
+ key = fun1.stateful_fun.get_arg_cache_key(x)
49
+ self.assertTrue(len(fun1.stateful_fun.get_states_by_cache(key)) == 2)
50
+
51
+ def test_kwargs(self):
52
+ a = bst.State(bst.random.randn(10))
53
+
54
+ @bst.compile.jit
55
+ def fun1(inp):
56
+ a.value += inp
57
+
58
+ b = bst.State(bst.random.randn(1))
59
+
60
+ def inner_fun(x):
61
+ b.value += x
62
+
63
+ bst.compile.for_loop(inner_fun, bst.random.randn(100))
64
+
65
+ return a.value + b.value
66
+
67
+ # test kwargs
68
+ print(fun1(inp=bst.random.randn(10)))
69
+
70
+ def test_jit_compile_sensitive_to_input_shape(self):
71
+ global_data = [0]
72
+
73
+ @bst.compile.jit
74
+ def fun1(inp):
75
+ global_data[0] += 1
76
+ return inp
77
+
78
+ print(fun1(1.))
79
+ self.assertTrue(global_data[0] == 1)
80
+
81
+ print(fun1(2.))
82
+ self.assertTrue(global_data[0] == 1)
83
+
84
+ print(fun1(bst.random.randn(10)))
85
+ self.assertTrue(global_data[0] == 2)
86
+
87
+ print(fun1(bst.random.randn(10, 10)))
88
+ self.assertTrue(global_data[0] == 3)
89
+
90
+ def test_jit_clear_cache(self):
91
+ a = bst.State(bst.random.randn(1))
92
+ compiling = []
93
+
94
+ @bst.compile.jit
95
+ def log2(x):
96
+ print('compiling')
97
+ compiling.append(1)
98
+ ln_x = jnp.log(x)
99
+ ln_2 = jnp.log(2.0) + a.value
100
+ return ln_x / ln_2
101
+
102
+ x = bst.random.randn(1)
103
+ print(log2(x)) # compiling
104
+ self.assertTrue(len(compiling) == 1)
105
+ print(log2(x)) # no compiling
106
+ self.assertTrue(len(compiling) == 1)
107
+
108
+ log2.clear_cache()
109
+ print(log2(x)) # compiling
110
+ self.assertTrue(len(compiling) == 2)
111
+
112
+ def test_jit_attribute_origin_fun(self):
113
+ def fun1(x):
114
+ return x
115
+
116
+ jitted_fun = bst.compile.jit(fun1)
117
+ self.assertTrue(jitted_fun.origin_fun is fun1)
118
+ self.assertTrue(isinstance(jitted_fun.stateful_fun, bst.compile.StatefulFunction))
119
+ self.assertTrue(callable(jitted_fun.jitted_fun))
120
+ self.assertTrue(callable(jitted_fun.clear_cache))
121
+
122
+ def test_clear_cache(self):
123
+ a = bst.State(bst.random.randn(1))
124
+
125
+ @bst.compile.jit
126
+ def f_jit(x, y):
127
+ print('Compiling')
128
+ a.value = jnp.sin(x) + jnp.cos(y)
129
+
130
+ f_jit(0.5, 1.0)
131
+ f_jit.clear_cache()
132
+ f_jit(0.5, 1.0)
133
+
134
+ def test_cache(self):
135
+ @bst.compile.jit
136
+ @bst.compile.jit
137
+ @bst.compile.jit
138
+ def f(a):
139
+ print('Compiling')
140
+ print(a)
141
+ return a + 1
142
+
143
+ print(f(1.))