brainstate 0.1.8__py2.py3-none-any.whl → 0.1.9__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 (133) hide show
  1. brainstate/__init__.py +58 -51
  2. brainstate/_compatible_import.py +148 -148
  3. brainstate/_state.py +1605 -1663
  4. brainstate/_state_test.py +52 -52
  5. brainstate/_utils.py +47 -47
  6. brainstate/augment/__init__.py +30 -30
  7. brainstate/augment/_autograd.py +778 -778
  8. brainstate/augment/_autograd_test.py +1289 -1289
  9. brainstate/augment/_eval_shape.py +99 -99
  10. brainstate/augment/_eval_shape_test.py +38 -38
  11. brainstate/augment/_mapping.py +1060 -1060
  12. brainstate/augment/_mapping_test.py +597 -597
  13. brainstate/augment/_random.py +151 -151
  14. brainstate/compile/__init__.py +38 -38
  15. brainstate/compile/_ad_checkpoint.py +204 -204
  16. brainstate/compile/_ad_checkpoint_test.py +49 -49
  17. brainstate/compile/_conditions.py +256 -256
  18. brainstate/compile/_conditions_test.py +220 -220
  19. brainstate/compile/_error_if.py +92 -92
  20. brainstate/compile/_error_if_test.py +52 -52
  21. brainstate/compile/_jit.py +346 -346
  22. brainstate/compile/_jit_test.py +143 -143
  23. brainstate/compile/_loop_collect_return.py +536 -536
  24. brainstate/compile/_loop_collect_return_test.py +58 -58
  25. brainstate/compile/_loop_no_collection.py +184 -184
  26. brainstate/compile/_loop_no_collection_test.py +50 -50
  27. brainstate/compile/_make_jaxpr.py +888 -888
  28. brainstate/compile/_make_jaxpr_test.py +156 -156
  29. brainstate/compile/_progress_bar.py +202 -202
  30. brainstate/compile/_unvmap.py +159 -159
  31. brainstate/compile/_util.py +147 -147
  32. brainstate/environ.py +563 -563
  33. brainstate/environ_test.py +62 -62
  34. brainstate/functional/__init__.py +27 -26
  35. brainstate/graph/__init__.py +29 -29
  36. brainstate/graph/_graph_node.py +244 -244
  37. brainstate/graph/_graph_node_test.py +73 -73
  38. brainstate/graph/_graph_operation.py +1738 -1738
  39. brainstate/graph/_graph_operation_test.py +563 -563
  40. brainstate/init/__init__.py +26 -26
  41. brainstate/init/_base.py +52 -52
  42. brainstate/init/_generic.py +244 -244
  43. brainstate/init/_random_inits.py +553 -553
  44. brainstate/init/_random_inits_test.py +149 -149
  45. brainstate/init/_regular_inits.py +105 -105
  46. brainstate/init/_regular_inits_test.py +50 -50
  47. brainstate/mixin.py +365 -363
  48. brainstate/mixin_test.py +77 -73
  49. brainstate/nn/__init__.py +135 -131
  50. brainstate/{functional → nn}/_activations.py +808 -813
  51. brainstate/{functional → nn}/_activations_test.py +331 -331
  52. brainstate/nn/_collective_ops.py +514 -514
  53. brainstate/nn/_collective_ops_test.py +43 -43
  54. brainstate/nn/_common.py +178 -178
  55. brainstate/nn/_conv.py +501 -501
  56. brainstate/nn/_conv_test.py +238 -238
  57. brainstate/nn/_delay.py +509 -502
  58. brainstate/nn/_delay_test.py +238 -184
  59. brainstate/nn/_dropout.py +426 -426
  60. brainstate/nn/_dropout_test.py +100 -100
  61. brainstate/nn/_dynamics.py +1343 -1343
  62. brainstate/nn/_dynamics_test.py +78 -78
  63. brainstate/nn/_elementwise.py +1119 -1119
  64. brainstate/nn/_elementwise_test.py +169 -169
  65. brainstate/nn/_embedding.py +58 -58
  66. brainstate/nn/_exp_euler.py +92 -92
  67. brainstate/nn/_exp_euler_test.py +35 -35
  68. brainstate/nn/_fixedprob.py +239 -239
  69. brainstate/nn/_fixedprob_test.py +114 -114
  70. brainstate/nn/_inputs.py +608 -608
  71. brainstate/nn/_linear.py +424 -424
  72. brainstate/nn/_linear_mv.py +83 -83
  73. brainstate/nn/_linear_mv_test.py +120 -120
  74. brainstate/nn/_linear_test.py +107 -107
  75. brainstate/nn/_ltp.py +28 -28
  76. brainstate/nn/_module.py +377 -377
  77. brainstate/nn/_module_test.py +40 -40
  78. brainstate/nn/_neuron.py +705 -705
  79. brainstate/nn/_neuron_test.py +161 -161
  80. brainstate/nn/_normalizations.py +975 -918
  81. brainstate/nn/_normalizations_test.py +73 -73
  82. brainstate/{functional → nn}/_others.py +46 -46
  83. brainstate/nn/_poolings.py +1177 -1177
  84. brainstate/nn/_poolings_test.py +217 -217
  85. brainstate/nn/_projection.py +486 -486
  86. brainstate/nn/_rate_rnns.py +554 -554
  87. brainstate/nn/_rate_rnns_test.py +63 -63
  88. brainstate/nn/_readout.py +209 -209
  89. brainstate/nn/_readout_test.py +53 -53
  90. brainstate/nn/_stp.py +236 -236
  91. brainstate/nn/_synapse.py +505 -505
  92. brainstate/nn/_synapse_test.py +131 -131
  93. brainstate/nn/_synaptic_projection.py +423 -423
  94. brainstate/nn/_synouts.py +162 -162
  95. brainstate/nn/_synouts_test.py +57 -57
  96. brainstate/nn/_utils.py +89 -89
  97. brainstate/nn/metrics.py +388 -388
  98. brainstate/optim/__init__.py +38 -38
  99. brainstate/optim/_base.py +64 -64
  100. brainstate/optim/_lr_scheduler.py +448 -448
  101. brainstate/optim/_lr_scheduler_test.py +50 -50
  102. brainstate/optim/_optax_optimizer.py +152 -152
  103. brainstate/optim/_optax_optimizer_test.py +53 -53
  104. brainstate/optim/_sgd_optimizer.py +1104 -1104
  105. brainstate/random/__init__.py +24 -24
  106. brainstate/random/_rand_funs.py +3616 -3616
  107. brainstate/random/_rand_funs_test.py +567 -567
  108. brainstate/random/_rand_seed.py +210 -210
  109. brainstate/random/_rand_seed_test.py +48 -48
  110. brainstate/random/_rand_state.py +1409 -1409
  111. brainstate/random/_random_for_unit.py +52 -52
  112. brainstate/surrogate.py +1957 -1957
  113. brainstate/transform.py +23 -23
  114. brainstate/typing.py +304 -304
  115. brainstate/util/__init__.py +50 -50
  116. brainstate/util/caller.py +98 -98
  117. brainstate/util/error.py +55 -55
  118. brainstate/util/filter.py +469 -469
  119. brainstate/util/others.py +540 -540
  120. brainstate/util/pretty_pytree.py +945 -945
  121. brainstate/util/pretty_pytree_test.py +159 -159
  122. brainstate/util/pretty_repr.py +328 -328
  123. brainstate/util/pretty_table.py +2954 -2954
  124. brainstate/util/scaling.py +258 -258
  125. brainstate/util/struct.py +523 -523
  126. {brainstate-0.1.8.dist-info → brainstate-0.1.9.dist-info}/METADATA +91 -99
  127. brainstate-0.1.9.dist-info/RECORD +130 -0
  128. {brainstate-0.1.8.dist-info → brainstate-0.1.9.dist-info}/WHEEL +1 -1
  129. {brainstate-0.1.8.dist-info → brainstate-0.1.9.dist-info/licenses}/LICENSE +202 -202
  130. brainstate/functional/_normalization.py +0 -81
  131. brainstate/functional/_spikes.py +0 -204
  132. brainstate-0.1.8.dist-info/RECORD +0 -132
  133. {brainstate-0.1.8.dist-info → brainstate-0.1.9.dist-info}/top_level.txt +0 -0
@@ -1,1289 +1,1289 @@
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
- from pprint import pprint
20
-
21
- import brainunit as u
22
- import jax
23
- import jax.numpy as jnp
24
- import pytest
25
-
26
- import brainstate
27
- from brainstate.augment._autograd import _jacfwd
28
-
29
-
30
- class TestPureFuncGrad(unittest.TestCase):
31
- def test_grad_pure_func_1(self):
32
- def call(a, b, c): return jnp.sum(a + b + c)
33
-
34
- brainstate.random.seed(1)
35
- a = jnp.ones(10)
36
- b = brainstate.random.randn(10)
37
- c = brainstate.random.uniform(size=10)
38
- f_grad = brainstate.augment.grad(call, argnums=[0, 1, 2])
39
- grads = f_grad(a, b, c)
40
-
41
- for g in grads: assert (g == 1.).all()
42
-
43
- def test_grad_pure_func_2(self):
44
- def call(a, b, c): return jnp.sum(a + b + c)
45
-
46
- brainstate.random.seed(1)
47
- a = jnp.ones(10)
48
- b = brainstate.random.randn(10)
49
- c = brainstate.random.uniform(size=10)
50
- f_grad = brainstate.augment.grad(call)
51
- assert (f_grad(a, b, c) == 1.).all()
52
-
53
- def test_grad_pure_func_aux1(self):
54
- def call(a, b, c):
55
- return jnp.sum(a + b + c), (jnp.sin(100), jnp.exp(0.1))
56
-
57
- brainstate.random.seed(1)
58
- f_grad = brainstate.augment.grad(call, argnums=[0, 1, 2])
59
- with pytest.raises(TypeError):
60
- f_grad(jnp.ones(10), brainstate.random.randn(10), brainstate.random.uniform(size=10))
61
-
62
- def test_grad_pure_func_aux2(self):
63
- def call(a, b, c):
64
- return jnp.sum(a + b + c), (jnp.sin(100), jnp.exp(0.1))
65
-
66
- brainstate.random.seed(1)
67
- f_grad = brainstate.augment.grad(call, argnums=[0, 1, 2], has_aux=True)
68
- grads, aux = f_grad(jnp.ones(10), brainstate.random.randn(10), brainstate.random.uniform(size=10))
69
- for g in grads: assert (g == 1.).all()
70
- assert aux[0] == jnp.sin(100)
71
- assert aux[1] == jnp.exp(0.1)
72
-
73
- def test_grad_pure_func_return1(self):
74
- def call(a, b, c): return jnp.sum(a + b + c)
75
-
76
- brainstate.random.seed(1)
77
- a = jnp.ones(10)
78
- b = brainstate.random.randn(10)
79
- c = brainstate.random.uniform(size=10)
80
- f_grad = brainstate.augment.grad(call, return_value=True)
81
- grads, returns = f_grad(a, b, c)
82
- assert (grads == 1.).all()
83
- assert returns == jnp.sum(a + b + c)
84
-
85
- def test_grad_func_return_aux1(self):
86
- def call(a, b, c):
87
- return jnp.sum(a + b + c), (jnp.sin(100), jnp.exp(0.1))
88
-
89
- brainstate.random.seed(1)
90
- a = jnp.ones(10)
91
- b = brainstate.random.randn(10)
92
- c = brainstate.random.uniform(size=10)
93
- f_grad = brainstate.augment.grad(call, return_value=True, has_aux=True)
94
- grads, returns, aux = f_grad(a, b, c)
95
- assert (grads == 1.).all()
96
- assert returns == jnp.sum(a + b + c)
97
- assert aux[0] == jnp.sin(100)
98
- assert aux[1] == jnp.exp(0.1)
99
-
100
-
101
- class TestObjectFuncGrad(unittest.TestCase):
102
- def test_grad_ob1(self):
103
- class Test(brainstate.nn.Module):
104
- def __init__(self):
105
- super(Test, self).__init__()
106
-
107
- self.a = brainstate.ParamState(jnp.ones(10))
108
- self.b = brainstate.ParamState(brainstate.random.randn(10))
109
- self.c = brainstate.ParamState(brainstate.random.uniform(size=10))
110
-
111
- def __call__(self):
112
- return jnp.sum(self.a.value + self.b.value + self.c.value)
113
-
114
- brainstate.random.seed(0)
115
-
116
- t = Test()
117
- f_grad = brainstate.augment.grad(t, grad_states={'a': t.a, 'b': t.b, 'c': t.c})
118
- grads = f_grad()
119
- for g in grads.values():
120
- assert (g == 1.).all()
121
-
122
- t = Test()
123
- f_grad = brainstate.augment.grad(t, grad_states=[t.a, t.b])
124
- grads = f_grad()
125
- for g in grads: assert (g == 1.).all()
126
-
127
- t = Test()
128
- f_grad = brainstate.augment.grad(t, grad_states=t.a)
129
- grads = f_grad()
130
- assert (grads == 1.).all()
131
-
132
- t = Test()
133
- f_grad = brainstate.augment.grad(t, grad_states=t.states())
134
- grads = f_grad()
135
- for g in grads.values():
136
- assert (g == 1.).all()
137
-
138
- def test_grad_ob_aux(self):
139
- class Test(brainstate.nn.Module):
140
- def __init__(self):
141
- super(Test, self).__init__()
142
- self.a = brainstate.ParamState(jnp.ones(10))
143
- self.b = brainstate.ParamState(brainstate.random.randn(10))
144
- self.c = brainstate.ParamState(brainstate.random.uniform(size=10))
145
-
146
- def __call__(self):
147
- return jnp.sum(self.a.value + self.b.value + self.c.value), (jnp.sin(100), jnp.exp(0.1))
148
-
149
- brainstate.random.seed(0)
150
- t = Test()
151
- f_grad = brainstate.augment.grad(t, grad_states=[t.a, t.b], has_aux=True)
152
- grads, aux = f_grad()
153
- for g in grads: assert (g == 1.).all()
154
- assert aux[0] == jnp.sin(100)
155
- assert aux[1] == jnp.exp(0.1)
156
-
157
- t = Test()
158
- f_grad = brainstate.augment.grad(t, grad_states=t.a, has_aux=True)
159
- grads, aux = f_grad()
160
- assert (grads == 1.).all()
161
- assert aux[0] == jnp.sin(100)
162
- assert aux[1] == jnp.exp(0.1)
163
-
164
- t = Test()
165
- f_grad = brainstate.augment.grad(t, grad_states=t.states(), has_aux=True)
166
- grads, aux = f_grad()
167
- self.assertTrue(len(grads) == len(t.states()))
168
-
169
- def test_grad_ob_return(self):
170
- class Test(brainstate.nn.Module):
171
- def __init__(self):
172
- super(Test, self).__init__()
173
- self.a = brainstate.ParamState(jnp.ones(10))
174
- self.b = brainstate.ParamState(brainstate.random.randn(10))
175
- self.c = brainstate.ParamState(brainstate.random.uniform(size=10))
176
-
177
- def __call__(self):
178
- return jnp.sum(self.a.value + self.b.value + self.c.value)
179
-
180
- brainstate.random.seed(0)
181
- t = Test()
182
- f_grad = brainstate.augment.grad(t, grad_states=[t.a, t.b], return_value=True)
183
- grads, returns = f_grad()
184
- for g in grads: assert (g == 1.).all()
185
- assert returns == t()
186
-
187
- t = Test()
188
- f_grad = brainstate.augment.grad(t, grad_states=t.a, return_value=True)
189
- grads, returns = f_grad()
190
- assert (grads == 1.).all()
191
- assert returns == t()
192
-
193
- def test_grad_ob_aux_return(self):
194
- class Test(brainstate.nn.Module):
195
- def __init__(self):
196
- super(Test, self).__init__()
197
- self.a = brainstate.ParamState(jnp.ones(10))
198
- self.b = brainstate.ParamState(brainstate.random.randn(10))
199
- self.c = brainstate.ParamState(brainstate.random.uniform(size=10))
200
-
201
- def __call__(self):
202
- return jnp.sum(self.a.value + self.b.value + self.c.value), (jnp.sin(100), jnp.exp(0.1))
203
-
204
- brainstate.random.seed(0)
205
- t = Test()
206
- f_grad = brainstate.augment.grad(t, grad_states=[t.a, t.b], has_aux=True, return_value=True)
207
- grads, returns, aux = f_grad()
208
- for g in grads: assert (g == 1.).all()
209
- assert returns == jnp.sum(t.a.value + t.b.value + t.c.value)
210
- assert aux[0] == jnp.sin(100)
211
- assert aux[1] == jnp.exp(0.1)
212
-
213
- t = Test()
214
- f_grad = brainstate.augment.grad(t, grad_states=t.a, has_aux=True, return_value=True)
215
- grads, returns, aux = f_grad()
216
- assert (grads == 1.).all()
217
- assert returns == jnp.sum(t.a.value + t.b.value + t.c.value)
218
- assert aux[0] == jnp.sin(100)
219
- assert aux[1] == jnp.exp(0.1)
220
-
221
- def test_grad_ob_argnums(self):
222
- class Test(brainstate.nn.Module):
223
- def __init__(self):
224
- super(Test, self).__init__()
225
- brainstate.random.seed()
226
- self.a = brainstate.ParamState(jnp.ones(10))
227
- self.b = brainstate.ParamState(brainstate.random.randn(10))
228
- self.c = brainstate.ParamState(brainstate.random.uniform(size=10))
229
-
230
- def __call__(self, d):
231
- return jnp.sum(self.a.value + self.b.value + self.c.value + 2 * d)
232
-
233
- brainstate.random.seed(0)
234
-
235
- t = Test()
236
- f_grad = brainstate.augment.grad(t, t.states(), argnums=0)
237
- var_grads, arg_grads = f_grad(brainstate.random.random(10))
238
- for g in var_grads.values(): assert (g == 1.).all()
239
- assert (arg_grads == 2.).all()
240
-
241
- t = Test()
242
- f_grad = brainstate.augment.grad(t, t.states(), argnums=[0])
243
- var_grads, arg_grads = f_grad(brainstate.random.random(10))
244
- for g in var_grads.values(): assert (g == 1.).all()
245
- assert (arg_grads[0] == 2.).all()
246
-
247
- t = Test()
248
- f_grad = brainstate.augment.grad(t, argnums=0)
249
- arg_grads = f_grad(brainstate.random.random(10))
250
- assert (arg_grads == 2.).all()
251
-
252
- t = Test()
253
- f_grad = brainstate.augment.grad(t, argnums=[0])
254
- arg_grads = f_grad(brainstate.random.random(10))
255
- assert (arg_grads[0] == 2.).all()
256
-
257
- def test_grad_ob_argnums_aux(self):
258
- class Test(brainstate.nn.Module):
259
- def __init__(self):
260
- super(Test, self).__init__()
261
- self.a = brainstate.ParamState(jnp.ones(10))
262
- self.b = brainstate.ParamState(brainstate.random.randn(10))
263
- self.c = brainstate.ParamState(brainstate.random.uniform(size=10))
264
-
265
- def __call__(self, d):
266
- return jnp.sum(self.a.value + self.b.value + self.c.value + 2 * d), (jnp.sin(100), jnp.exp(0.1))
267
-
268
- brainstate.random.seed(0)
269
-
270
- t = Test()
271
- f_grad = brainstate.augment.grad(t, grad_states=t.states(), argnums=0, has_aux=True)
272
- (var_grads, arg_grads), aux = f_grad(brainstate.random.random(10))
273
- for g in var_grads.values(): assert (g == 1.).all()
274
- assert (arg_grads == 2.).all()
275
- assert aux[0] == jnp.sin(100)
276
- assert aux[1] == jnp.exp(0.1)
277
-
278
- t = Test()
279
- f_grad = brainstate.augment.grad(t, grad_states=t.states(), argnums=[0], has_aux=True)
280
- (var_grads, arg_grads), aux = f_grad(brainstate.random.random(10))
281
- for g in var_grads.values(): assert (g == 1.).all()
282
- assert (arg_grads[0] == 2.).all()
283
- assert aux[0] == jnp.sin(100)
284
- assert aux[1] == jnp.exp(0.1)
285
-
286
- t = Test()
287
- f_grad = brainstate.augment.grad(t, argnums=0, has_aux=True)
288
- arg_grads, aux = f_grad(brainstate.random.random(10))
289
- assert (arg_grads == 2.).all()
290
- assert aux[0] == jnp.sin(100)
291
- assert aux[1] == jnp.exp(0.1)
292
-
293
- t = Test()
294
- f_grad = brainstate.augment.grad(t, argnums=[0], has_aux=True)
295
- arg_grads, aux = f_grad(brainstate.random.random(10))
296
- assert (arg_grads[0] == 2.).all()
297
- assert aux[0] == jnp.sin(100)
298
- assert aux[1] == jnp.exp(0.1)
299
-
300
- def test_grad_ob_argnums_return(self):
301
- class Test(brainstate.nn.Module):
302
- def __init__(self):
303
- super(Test, self).__init__()
304
-
305
- self.a = brainstate.ParamState(jnp.ones(10))
306
- self.b = brainstate.ParamState(brainstate.random.randn(10))
307
- self.c = brainstate.ParamState(brainstate.random.uniform(size=10))
308
-
309
- def __call__(self, d):
310
- return jnp.sum(self.a.value + self.b.value + self.c.value + 2 * d)
311
-
312
- brainstate.random.seed(0)
313
-
314
- t = Test()
315
- f_grad = brainstate.augment.grad(t, t.states(), argnums=0, return_value=True)
316
- d = brainstate.random.random(10)
317
- (var_grads, arg_grads), loss = f_grad(d)
318
- for g in var_grads.values():
319
- assert (g == 1.).all()
320
- assert (arg_grads == 2.).all()
321
- assert loss == t(d)
322
-
323
- t = Test()
324
- f_grad = brainstate.augment.grad(t, t.states(), argnums=[0], return_value=True)
325
- d = brainstate.random.random(10)
326
- (var_grads, arg_grads), loss = f_grad(d)
327
- for g in var_grads.values():
328
- assert (g == 1.).all()
329
- assert (arg_grads[0] == 2.).all()
330
- assert loss == t(d)
331
-
332
- t = Test()
333
- f_grad = brainstate.augment.grad(t, argnums=0, return_value=True)
334
- d = brainstate.random.random(10)
335
- arg_grads, loss = f_grad(d)
336
- assert (arg_grads == 2.).all()
337
- assert loss == t(d)
338
-
339
- t = Test()
340
- f_grad = brainstate.augment.grad(t, argnums=[0], return_value=True)
341
- d = brainstate.random.random(10)
342
- arg_grads, loss = f_grad(d)
343
- assert (arg_grads[0] == 2.).all()
344
- assert loss == t(d)
345
-
346
- def test_grad_ob_argnums_aux_return(self):
347
- class Test(brainstate.nn.Module):
348
- def __init__(self):
349
- super(Test, self).__init__()
350
- self.a = brainstate.ParamState(jnp.ones(10))
351
- self.b = brainstate.ParamState(brainstate.random.randn(10))
352
- self.c = brainstate.ParamState(brainstate.random.uniform(size=10))
353
-
354
- def __call__(self, d):
355
- return jnp.sum(self.a.value + self.b.value + self.c.value + 2 * d), (jnp.sin(100), jnp.exp(0.1))
356
-
357
- brainstate.random.seed(0)
358
-
359
- t = Test()
360
- f_grad = brainstate.augment.grad(t, grad_states=t.states(), argnums=0, has_aux=True, return_value=True)
361
- d = brainstate.random.random(10)
362
- (var_grads, arg_grads), loss, aux = f_grad(d)
363
- for g in var_grads.values(): assert (g == 1.).all()
364
- assert (arg_grads == 2.).all()
365
- assert aux[0] == jnp.sin(100)
366
- assert aux[1] == jnp.exp(0.1)
367
- assert loss == t(d)[0]
368
-
369
- t = Test()
370
- f_grad = brainstate.augment.grad(t, grad_states=t.states(), argnums=[0], has_aux=True, return_value=True)
371
- d = brainstate.random.random(10)
372
- (var_grads, arg_grads), loss, aux = f_grad(d)
373
- for g in var_grads.values(): assert (g == 1.).all()
374
- assert (arg_grads[0] == 2.).all()
375
- assert aux[0] == jnp.sin(100)
376
- assert aux[1] == jnp.exp(0.1)
377
- assert loss == t(d)[0]
378
-
379
- t = Test()
380
- f_grad = brainstate.augment.grad(t, argnums=0, has_aux=True, return_value=True)
381
- d = brainstate.random.random(10)
382
- arg_grads, loss, aux = f_grad(d)
383
- assert (arg_grads == 2.).all()
384
- assert aux[0] == jnp.sin(100)
385
- assert aux[1] == jnp.exp(0.1)
386
- assert loss == t(d)[0]
387
-
388
- t = Test()
389
- f_grad = brainstate.augment.grad(t, argnums=[0], has_aux=True, return_value=True)
390
- d = brainstate.random.random(10)
391
- arg_grads, loss, aux = f_grad(d)
392
- assert (arg_grads[0] == 2.).all()
393
- assert aux[0] == jnp.sin(100)
394
- assert aux[1] == jnp.exp(0.1)
395
- assert loss == t(d)[0]
396
-
397
-
398
- class TestPureFuncJacobian(unittest.TestCase):
399
- def test1(self):
400
- jac, aux = _jacfwd(lambda x: (x ** 3, [x ** 2]), has_aux=True)(3.)
401
- self.assertTrue(jax.numpy.allclose(jac, jax.jacfwd(lambda x: x ** 3)(3.)))
402
- self.assertTrue(aux[0] == 9.)
403
-
404
- def test_jacfwd_and_aux_nested(self):
405
- def f(x):
406
- jac, aux = _jacfwd(lambda x: (x ** 3, [x ** 3]), has_aux=True)(x)
407
- return aux[0]
408
-
409
- f2 = lambda x: x ** 3
410
-
411
- self.assertEqual(_jacfwd(f)(4.), _jacfwd(f2)(4.))
412
- self.assertEqual(jax.jit(_jacfwd(f))(4.), _jacfwd(f2)(4.))
413
- self.assertEqual(jax.jit(_jacfwd(jax.jit(f)))(4.), _jacfwd(f2)(4.))
414
-
415
- self.assertEqual(_jacfwd(f)(jnp.asarray(4.)), _jacfwd(f2)(jnp.asarray(4.)))
416
- self.assertEqual(jax.jit(_jacfwd(f))(jnp.asarray(4.)), _jacfwd(f2)(jnp.asarray(4.)))
417
- self.assertEqual(jax.jit(_jacfwd(jax.jit(f)))(jnp.asarray(4.)), _jacfwd(f2)(jnp.asarray(4.)))
418
-
419
- def f(x):
420
- jac, aux = _jacfwd(lambda x: (x ** 3, [x ** 3]), has_aux=True)(x)
421
- return aux[0] * jnp.sin(x)
422
-
423
- f2 = lambda x: x ** 3 * jnp.sin(x)
424
-
425
- self.assertEqual(_jacfwd(f)(4.), _jacfwd(f2)(4.))
426
- self.assertEqual(jax.jit(_jacfwd(f))(4.), _jacfwd(f2)(4.))
427
- self.assertEqual(jax.jit(_jacfwd(jax.jit(f)))(4.), _jacfwd(f2)(4.))
428
-
429
- self.assertEqual(_jacfwd(f)(jnp.asarray(4.)), _jacfwd(f2)(jnp.asarray(4.)))
430
- self.assertEqual(jax.jit(_jacfwd(f))(jnp.asarray(4.)), _jacfwd(f2)(jnp.asarray(4.)))
431
- self.assertEqual(jax.jit(_jacfwd(jax.jit(f)))(jnp.asarray(4.)), _jacfwd(f2)(jnp.asarray(4.)))
432
-
433
- def test_jacrev1(self):
434
- def f1(x, y):
435
- r = jnp.asarray([x[0] * y[0], 5 * x[2] * y[1], 4 * x[1] ** 2 - 2 * x[2], x[2] * jnp.sin(x[0])])
436
- return r
437
-
438
- br = brainstate.augment.jacrev(f1)(jnp.array([1., 2., 3.]), jnp.array([10., 5.]))
439
- jr = jax.jacrev(f1)(jnp.array([1., 2., 3.]), jnp.array([10., 5.]))
440
- assert (br == jr).all()
441
-
442
- br = brainstate.augment.jacrev(f1, argnums=(0, 1))(jnp.array([1., 2., 3.]), jnp.array([10., 5.]))
443
- jr = jax.jacrev(f1, argnums=(0, 1))(jnp.array([1., 2., 3.]), jnp.array([10., 5.]))
444
- assert (br[0] == jr[0]).all()
445
- assert (br[1] == jr[1]).all()
446
-
447
- def test_jacrev2(self):
448
- print()
449
-
450
- def f2(x, y):
451
- r1 = jnp.asarray([x[0] * y[0], 5 * x[2] * y[1]])
452
- r2 = jnp.asarray([4 * x[1] ** 2 - 2 * x[2], x[2] * jnp.sin(x[0])])
453
- return r1, r2
454
-
455
- jr = jax.jacrev(f2)(jnp.array([1., 2., 3.]), jnp.array([10., 5.]))
456
- pprint(jr)
457
-
458
- br = brainstate.augment.jacrev(f2)(jnp.array([1., 2., 3.]), jnp.array([10., 5.]))
459
- pprint(br)
460
- assert jnp.array_equal(br[0], jr[0])
461
- assert jnp.array_equal(br[1], jr[1])
462
-
463
- br = brainstate.augment.jacrev(f2)(jnp.array([1., 2., 3.]), jnp.array([10., 5.]))
464
- pprint(br)
465
- assert jnp.array_equal(br[0], jr[0])
466
- assert jnp.array_equal(br[1], jr[1])
467
-
468
- def f2(x, y):
469
- r1 = jnp.asarray([x[0] * y[0], 5 * x[2] * y[1]])
470
- r2 = jnp.asarray([4 * x[1] ** 2 - 2 * x[2], x[2] * jnp.sin(x[0])])
471
- return r1, r2
472
-
473
- br = brainstate.augment.jacrev(f2)(jnp.array([1., 2., 3.]), jnp.array([10., 5.]))
474
- pprint(br)
475
- assert jnp.array_equal(br[0], jr[0])
476
- assert jnp.array_equal(br[1], jr[1])
477
-
478
- br = brainstate.augment.jacrev(f2)(jnp.array([1., 2., 3.]), jnp.array([10., 5.]))
479
- pprint(br)
480
- assert jnp.array_equal(br[0], jr[0])
481
- assert jnp.array_equal(br[1], jr[1])
482
-
483
- def test_jacrev3(self):
484
- print()
485
-
486
- def f3(x, y):
487
- r1 = jnp.asarray([x[0] * y[0], 5 * x[2] * y[1]])
488
- r2 = jnp.asarray([4 * x[1] ** 2 - 2 * x[2], x[2] * jnp.sin(x[0])])
489
- return r1, r2
490
-
491
- jr = jax.jacrev(f3, argnums=(0, 1))(jnp.array([1., 2., 3.]), jnp.array([10., 5.]))
492
- pprint(jr)
493
-
494
- br = brainstate.augment.jacrev(f3, argnums=(0, 1))(jnp.array([1., 2., 3.]), jnp.array([10., 5.]))
495
- pprint(br)
496
- assert jnp.array_equal(br[0][0], jr[0][0])
497
- assert jnp.array_equal(br[0][1], jr[0][1])
498
- assert jnp.array_equal(br[1][0], jr[1][0])
499
- assert jnp.array_equal(br[1][1], jr[1][1])
500
-
501
- br = brainstate.augment.jacrev(f3, argnums=(0, 1))(jnp.array([1., 2., 3.]), jnp.array([10., 5.]))
502
- pprint(br)
503
- assert jnp.array_equal(br[0][0], jr[0][0])
504
- assert jnp.array_equal(br[0][1], jr[0][1])
505
- assert jnp.array_equal(br[1][0], jr[1][0])
506
- assert jnp.array_equal(br[1][1], jr[1][1])
507
-
508
- def f3(x, y):
509
- r1 = jnp.asarray([x[0] * y[0], 5 * x[2] * y[1]])
510
- r2 = jnp.asarray([4 * x[1] ** 2 - 2 * x[2], x[2] * jnp.sin(x[0])])
511
- return r1, r2
512
-
513
- br = brainstate.augment.jacrev(f3, argnums=(0, 1))(jnp.array([1., 2., 3.]), jnp.array([10., 5.]))
514
- pprint(br)
515
- assert jnp.array_equal(br[0][0], jr[0][0])
516
- assert jnp.array_equal(br[0][1], jr[0][1])
517
- assert jnp.array_equal(br[1][0], jr[1][0])
518
- assert jnp.array_equal(br[1][1], jr[1][1])
519
-
520
- br = brainstate.augment.jacrev(f3, argnums=(0, 1))(jnp.array([1., 2., 3.]), jnp.array([10., 5.]))
521
- pprint(br)
522
- assert jnp.array_equal(br[0][0], jr[0][0])
523
- assert jnp.array_equal(br[0][1], jr[0][1])
524
- assert jnp.array_equal(br[1][0], jr[1][0])
525
- assert jnp.array_equal(br[1][1], jr[1][1])
526
-
527
- def test_jacrev_aux1(self):
528
- x = jnp.array([1., 2., 3.])
529
- y = jnp.array([10., 5.])
530
-
531
- def f1(x, y):
532
- a = 4 * x[1] ** 2 - 2 * x[2]
533
- r = jnp.asarray([x[0] * y[0], 5 * x[2] * y[1], a, x[2] * jnp.sin(x[0])])
534
- return r, a
535
-
536
- f2 = lambda *args: f1(*args)[0]
537
- jr = jax.jacrev(f2)(x, y) # jax jacobian
538
- pprint(jr)
539
- grads, aux = brainstate.augment.jacrev(f1, has_aux=True)(x, y)
540
- assert (grads == jr).all()
541
- assert aux == (4 * x[1] ** 2 - 2 * x[2])
542
-
543
- jr = jax.jacrev(f2, argnums=(0, 1))(x, y) # jax jacobian
544
- pprint(jr)
545
- grads, aux = brainstate.augment.jacrev(f1, argnums=(0, 1), has_aux=True)(x, y)
546
- assert (grads[0] == jr[0]).all()
547
- assert (grads[1] == jr[1]).all()
548
- assert aux == (4 * x[1] ** 2 - 2 * x[2])
549
-
550
- def test_jacrev_return_aux1(self):
551
- with brainstate.environ.context(precision=64):
552
- def f1(x, y):
553
- a = 4 * x[1] ** 2 - 2 * x[2]
554
- r = jnp.asarray([x[0] * y[0], 5 * x[2] * y[1], a, x[2] * jnp.sin(x[0])])
555
- return r, a
556
-
557
- _x = jnp.array([1., 2., 3.])
558
- _y = jnp.array([10., 5.])
559
- _r, _a = f1(_x, _y)
560
- f2 = lambda *args: f1(*args)[0]
561
- _g1 = jax.jacrev(f2)(_x, _y) # jax jacobian
562
- pprint(_g1)
563
- _g2 = jax.jacrev(f2, argnums=(0, 1))(_x, _y) # jax jacobian
564
- pprint(_g2)
565
-
566
- grads, vec, aux = brainstate.augment.jacrev(f1, return_value=True, has_aux=True)(_x, _y)
567
- assert (grads == _g1).all()
568
- assert aux == _a
569
- assert (vec == _r).all()
570
-
571
- grads, vec, aux = brainstate.augment.jacrev(f1, return_value=True, argnums=(0, 1), has_aux=True)(_x, _y)
572
- assert (grads[0] == _g2[0]).all()
573
- assert (grads[1] == _g2[1]).all()
574
- assert aux == _a
575
- assert (vec == _r).all()
576
-
577
-
578
- class TestClassFuncJacobian(unittest.TestCase):
579
- def test_jacrev1(self):
580
- def f1(x, y):
581
- r = jnp.asarray([x[0] * y[0], 5 * x[2] * y[1], 4 * x[1] ** 2 - 2 * x[2], x[2] * jnp.sin(x[0])])
582
- return r
583
-
584
- _x = jnp.array([1., 2., 3.])
585
- _y = jnp.array([10., 5.])
586
-
587
- class Test(brainstate.nn.Module):
588
- def __init__(self):
589
- super(Test, self).__init__()
590
- self.x = brainstate.State(jnp.array([1., 2., 3.]))
591
- self.y = brainstate.State(jnp.array([10., 5.]))
592
-
593
- def __call__(self, ):
594
- a = self.x.value[0] * self.y.value[0]
595
- b = 5 * self.x.value[2] * self.y.value[1]
596
- c = 4 * self.x.value[1] ** 2 - 2 * self.x.value[2]
597
- d = self.x.value[2] * jnp.sin(self.x.value[0])
598
- r = jnp.asarray([a, b, c, d])
599
- return r
600
-
601
- _jr = jax.jacrev(f1)(_x, _y)
602
- t = Test()
603
- br = brainstate.augment.jacrev(t, grad_states=t.x)()
604
- self.assertTrue((br == _jr).all())
605
-
606
- _jr = jax.jacrev(f1, argnums=(0, 1))(_x, _y)
607
- t = Test()
608
- br = brainstate.augment.jacrev(t, grad_states=[t.x, t.y])()
609
- self.assertTrue((br[0] == _jr[0]).all())
610
- self.assertTrue((br[1] == _jr[1]).all())
611
-
612
-
613
- #
614
- # def test_jacfwd1(self):
615
- # def f1(x, y):
616
- # r = jnp.asarray([x[0] * y[0], 5 * x[2] * y[1], 4 * x[1] ** 2 - 2 * x[2], x[2] * jnp.sin(x[0])])
617
- # return r
618
- #
619
- # _x = jnp.array([1., 2., 3.])
620
- # _y = jnp.array([10., 5.])
621
- #
622
- # class Test(brainstate.nn.Module):
623
- # def __init__(self):
624
- # super(Test, self).__init__()
625
- # self.x = jnp.Variable(jnp.array([1., 2., 3.]))
626
- # self.y = jnp.Variable(jnp.array([10., 5.]))
627
- #
628
- # def __call__(self, ):
629
- # a = self.x[0] * self.y[0]
630
- # b = 5 * self.x[2] * self.y[1]
631
- # c = 4 * self.x[1] ** 2 - 2 * self.x[2]
632
- # d = self.x[2] * jnp.sin(self.x[0])
633
- # r = jnp.asarray([a, b, c, d])
634
- # return r
635
- #
636
- # _jr = jax.jacfwd(f1)(_x, _y)
637
- # t = Test()
638
- # br = brainstate.augment.jacfwd(t, grad_states=t.x)()
639
- # self.assertTrue((br == _jr).all())
640
- #
641
- # _jr = jax.jacfwd(f1, argnums=(0, 1))(_x, _y)
642
- # t = Test()
643
- # br = brainstate.augment.jacfwd(t, grad_states=[t.x, t.y])()
644
- # self.assertTrue((br[0] == _jr[0]).all())
645
- # self.assertTrue((br[1] == _jr[1]).all())
646
- #
647
- # def test_jacrev2(self):
648
- # def f1(x, y):
649
- # r = jnp.asarray([x[0] * y[0], 5 * x[2] * y[1], 4 * x[1] ** 2 - 2 * x[2], x[2] * jnp.sin(x[0])])
650
- # return r
651
- #
652
- # _x = jnp.array([1., 2., 3.])
653
- # _y = jnp.array([10., 5.])
654
- #
655
- # class Test(brainstate.nn.Module):
656
- # def __init__(self):
657
- # super(Test, self).__init__()
658
- # self.x = jnp.Variable(jnp.array([1., 2., 3.]))
659
- #
660
- # def __call__(self, y):
661
- # a = self.x[0] * y[0]
662
- # b = 5 * self.x[2] * y[1]
663
- # c = 4 * self.x[1] ** 2 - 2 * self.x[2]
664
- # d = self.x[2] * jnp.sin(self.x[0])
665
- # r = jnp.asarray([a, b, c, d])
666
- # return r
667
- #
668
- # _jr = jax.jacrev(f1)(_x, _y)
669
- # t = Test()
670
- # br = brainstate.augment.jacrev(t, grad_states=t.x)(_y)
671
- # self.assertTrue((br == _jr).all())
672
- #
673
- # _jr = jax.jacrev(f1, argnums=(0, 1))(_x, _y)
674
- # t = Test()
675
- # var_grads, arg_grads = brainstate.augment.jacrev(t, grad_states=t.x, argnums=0)(_y)
676
- # print(var_grads, )
677
- # print(arg_grads, )
678
- # self.assertTrue((var_grads == _jr[0]).all())
679
- # self.assertTrue((arg_grads == _jr[1]).all())
680
- #
681
- # def test_jacfwd2(self):
682
- # def f1(x, y):
683
- # r = jnp.asarray([x[0] * y[0], 5 * x[2] * y[1], 4 * x[1] ** 2 - 2 * x[2], x[2] * jnp.sin(x[0])])
684
- # return r
685
- #
686
- # _x = jnp.array([1., 2., 3.])
687
- # _y = jnp.array([10., 5.])
688
- #
689
- # class Test(brainstate.nn.Module):
690
- # def __init__(self):
691
- # super(Test, self).__init__()
692
- # self.x = jnp.Variable(jnp.array([1., 2., 3.]))
693
- #
694
- # def __call__(self, y):
695
- # a = self.x[0] * y[0]
696
- # b = 5 * self.x[2] * y[1]
697
- # c = 4 * self.x[1] ** 2 - 2 * self.x[2]
698
- # d = self.x[2] * jnp.sin(self.x[0])
699
- # r = jnp.asarray([a, b, c, d])
700
- # return r
701
- #
702
- # _jr = jax.jacfwd(f1)(_x, _y)
703
- # t = Test()
704
- # br = brainstate.augment.jacfwd(t, grad_states=t.x)(_y)
705
- # self.assertTrue((br == _jr).all())
706
- #
707
- # _jr = jax.jacfwd(f1, argnums=(0, 1))(_x, _y)
708
- # t = Test()
709
- # var_grads, arg_grads = brainstate.augment.jacfwd(t, grad_states=t.x, argnums=0)(_y)
710
- # print(var_grads, )
711
- # print(arg_grads, )
712
- # self.assertTrue((var_grads == _jr[0]).all())
713
- # self.assertTrue((arg_grads == _jr[1]).all())
714
- #
715
- # def test_jacrev_aux1(self):
716
- # jnp.enable_x64()
717
- #
718
- # def f1(x, y):
719
- # r = jnp.asarray([x[0] * y[0], 5 * x[2] * y[1], 4 * x[1] ** 2 - 2 * x[2], x[2] * jnp.sin(x[0])])
720
- # return r
721
- #
722
- # _x = jnp.array([1., 2., 3.])
723
- # _y = jnp.array([10., 5.])
724
- #
725
- # class Test(brainstate.nn.Module):
726
- # def __init__(self):
727
- # super(Test, self).__init__()
728
- # self.x = jnp.Variable(jnp.array([1., 2., 3.]))
729
- #
730
- # def __call__(self, y):
731
- # a = self.x[0] * y[0]
732
- # b = 5 * self.x[2] * y[1]
733
- # c = 4 * self.x[1] ** 2 - 2 * self.x[2]
734
- # d = self.x[2] * jnp.sin(self.x[0])
735
- # r = jnp.asarray([a, b, c, d])
736
- # return r, (c, d)
737
- #
738
- # _jr = jax.jacrev(f1)(_x, _y)
739
- # t = Test()
740
- # br, _ = brainstate.augment.jacrev(t, grad_states=t.x, has_aux=True)(_y)
741
- # self.assertTrue((br == _jr).all())
742
- #
743
- # t = Test()
744
- # _jr = jax.jacrev(f1, argnums=(0, 1))(_x, _y)
745
- # _aux = t(_y)[1]
746
- # (var_grads, arg_grads), aux = brainstate.augment.jacrev(t, grad_states=t.x, argnums=0, has_aux=True)(_y)
747
- # print(var_grads, )
748
- # print(arg_grads, )
749
- # self.assertTrue((var_grads == _jr[0]).all())
750
- # self.assertTrue((arg_grads == _jr[1]).all())
751
- # self.assertTrue(jnp.array_equal(aux, _aux))
752
- #
753
- # jnp.disable_x64()
754
- #
755
- # def test_jacfwd_aux1(self):
756
- # jnp.enable_x64()
757
- #
758
- # def f1(x, y):
759
- # r = jnp.asarray([x[0] * y[0], 5 * x[2] * y[1], 4 * x[1] ** 2 - 2 * x[2], x[2] * jnp.sin(x[0])])
760
- # return r
761
- #
762
- # _x = jnp.array([1., 2., 3.])
763
- # _y = jnp.array([10., 5.])
764
- #
765
- # class Test(brainstate.nn.Module):
766
- # def __init__(self):
767
- # super(Test, self).__init__()
768
- # self.x = jnp.Variable(jnp.array([1., 2., 3.]))
769
- #
770
- # def __call__(self, y):
771
- # a = self.x[0] * y[0]
772
- # b = 5 * self.x[2] * y[1]
773
- # c = 4 * self.x[1] ** 2 - 2 * self.x[2]
774
- # d = self.x[2] * jnp.sin(self.x[0])
775
- # r = jnp.asarray([a, b, c, d])
776
- # return r, (c, d)
777
- #
778
- # _jr = jax.jacfwd(f1)(_x, _y)
779
- # t = Test()
780
- # br, (c, d) = brainstate.augment.jacfwd(t, grad_states=t.x, has_aux=True)(_y)
781
- # # print(_jr)
782
- # # print(br)
783
- # a = (br == _jr)
784
- # self.assertTrue(a.all())
785
- #
786
- # t = Test()
787
- # _jr = jax.jacfwd(f1, argnums=(0, 1))(_x, _y)
788
- # _aux = t(_y)[1]
789
- # (var_grads, arg_grads), aux = brainstate.augment.jacfwd(t, grad_states=t.x, argnums=0, has_aux=True)(_y)
790
- # print(var_grads, )
791
- # print(arg_grads, )
792
- # self.assertTrue((var_grads == _jr[0]).all())
793
- # self.assertTrue((arg_grads == _jr[1]).all())
794
- # self.assertTrue(jnp.array_equal(aux, _aux))
795
- #
796
- # jnp.disable_x64()
797
- #
798
- # def test_jacrev_return_aux1(self):
799
- # jnp.enable_x64()
800
- #
801
- # def f1(x, y):
802
- # r = jnp.asarray([x[0] * y[0], 5 * x[2] * y[1], 4 * x[1] ** 2 - 2 * x[2], x[2] * jnp.sin(x[0])])
803
- # return r
804
- #
805
- # _x = jnp.array([1., 2., 3.])
806
- # _y = jnp.array([10., 5.])
807
- #
808
- # class Test(brainstate.nn.Module):
809
- # def __init__(self):
810
- # super(Test, self).__init__()
811
- # self.x = jnp.Variable(jnp.array([1., 2., 3.]))
812
- #
813
- # def __call__(self, y):
814
- # a = self.x[0] * y[0]
815
- # b = 5 * self.x[2] * y[1]
816
- # c = 4 * self.x[1] ** 2 - 2 * self.x[2]
817
- # d = self.x[2] * jnp.sin(self.x[0])
818
- # r = jnp.asarray([a, b, c, d])
819
- # return r, (c, d)
820
- #
821
- # _jr = jax.jacrev(f1)(_x, _y)
822
- # t = Test()
823
- # br, _ = brainstate.augment.jacrev(t, grad_states=t.x, has_aux=True)(_y)
824
- # self.assertTrue((br == _jr).all())
825
- #
826
- # t = Test()
827
- # _jr = jax.jacrev(f1, argnums=(0, 1))(_x, _y)
828
- # _val, _aux = t(_y)
829
- # (var_grads, arg_grads), value, aux = brainstate.augment.jacrev(t, grad_states=t.x, argnums=0, has_aux=True, return_value=True)(_y)
830
- # print(var_grads, )
831
- # print(arg_grads, )
832
- # self.assertTrue((var_grads == _jr[0]).all())
833
- # self.assertTrue((arg_grads == _jr[1]).all())
834
- # self.assertTrue(jnp.array_equal(aux, _aux))
835
- # self.assertTrue(jnp.array_equal(value, _val))
836
- #
837
- # jnp.disable_x64()
838
- #
839
- # def test_jacfwd_return_aux1(self):
840
- # jnp.enable_x64()
841
- #
842
- # def f1(x, y):
843
- # r = jnp.asarray([x[0] * y[0], 5 * x[2] * y[1], 4 * x[1] ** 2 - 2 * x[2], x[2] * jnp.sin(x[0])])
844
- # return r
845
- #
846
- # _x = jnp.array([1., 2., 3.])
847
- # _y = jnp.array([10., 5.])
848
- #
849
- # class Test(brainstate.nn.Module):
850
- # def __init__(self):
851
- # super(Test, self).__init__()
852
- # self.x = jnp.Variable(jnp.array([1., 2., 3.]))
853
- #
854
- # def __call__(self, y):
855
- # a = self.x[0] * y[0]
856
- # b = 5 * self.x[2] * y[1]
857
- # c = 4 * self.x[1] ** 2 - 2 * self.x[2]
858
- # d = self.x[2] * jnp.sin(self.x[0])
859
- # r = jnp.asarray([a, b, c, d])
860
- # return r, (c, d)
861
- #
862
- # _jr = jax.jacfwd(f1)(_x, _y)
863
- # t = Test()
864
- # br, _ = brainstate.augment.jacfwd(t, grad_states=t.x, has_aux=True)(_y)
865
- # self.assertTrue((br == _jr).all())
866
- #
867
- # t = Test()
868
- # _jr = jax.jacfwd(f1, argnums=(0, 1))(_x, _y)
869
- # _val, _aux = t(_y)
870
- # (var_grads, arg_grads), value, aux = brainstate.augment.jacfwd(t, grad_states=t.x, argnums=0, has_aux=True, return_value=True)(_y)
871
- # print(_val, )
872
- # print('_aux: ', _aux, 'aux: ', aux)
873
- # print(var_grads, )
874
- # print(arg_grads, )
875
- # self.assertTrue((var_grads == _jr[0]).all())
876
- # self.assertTrue((arg_grads == _jr[1]).all())
877
- # self.assertTrue(jnp.array_equal(aux, _aux))
878
- # self.assertTrue(jnp.array_equal(value, _val))
879
- #
880
- # jnp.disable_x64()
881
- #
882
- #
883
- # class TestPureFuncVectorGrad(unittest.TestCase):
884
- # def test1(self):
885
- # f = lambda x: 3 * x ** 2
886
- # _x = jnp.ones(10)
887
- # pprint(brainstate.augment.vector_grad(f, argnums=0)(_x))
888
- #
889
- # def test2(self):
890
- # def f(x, y):
891
- # dx = x ** 2 + y ** 2 + 10
892
- # return dx
893
- #
894
- # _x = jnp.ones(5)
895
- # _y = jnp.ones(5)
896
- #
897
- # g = brainstate.augment.vector_grad(f, argnums=0)(_x, _y)
898
- # pprint(g)
899
- # self.assertTrue(jnp.array_equal(g, 2 * _x))
900
- #
901
- # g = brainstate.augment.vector_grad(f, argnums=(0,))(_x, _y)
902
- # self.assertTrue(jnp.array_equal(g[0], 2 * _x))
903
- #
904
- # g = brainstate.augment.vector_grad(f, argnums=(0, 1))(_x, _y)
905
- # pprint(g)
906
- # self.assertTrue(jnp.array_equal(g[0], 2 * _x))
907
- # self.assertTrue(jnp.array_equal(g[1], 2 * _y))
908
- #
909
- # def test3(self):
910
- # def f(x, y):
911
- # dx = x ** 2 + y ** 2 + 10
912
- # dy = x ** 3 + y ** 3 - 10
913
- # return dx, dy
914
- #
915
- # _x = jnp.ones(5)
916
- # _y = jnp.ones(5)
917
- #
918
- # g = brainstate.augment.vector_grad(f, argnums=0)(_x, _y)
919
- # # pprint(g)
920
- # self.assertTrue(jnp.array_equal(g, 2 * _x + 3 * _x ** 2))
921
- #
922
- # g = brainstate.augment.vector_grad(f, argnums=(0,))(_x, _y)
923
- # self.assertTrue(jnp.array_equal(g[0], 2 * _x + 3 * _x ** 2))
924
- #
925
- # g = brainstate.augment.vector_grad(f, argnums=(0, 1))(_x, _y)
926
- # # pprint(g)
927
- # self.assertTrue(jnp.array_equal(g[0], 2 * _x + 3 * _x ** 2))
928
- # self.assertTrue(jnp.array_equal(g[1], 2 * _y + 3 * _y ** 2))
929
- #
930
- # def test4_2d(self):
931
- # def f(x, y):
932
- # dx = x ** 2 + y ** 2 + 10
933
- # return dx
934
- #
935
- # _x = jnp.ones((5, 5))
936
- # _y = jnp.ones((5, 5))
937
- #
938
- # g = brainstate.augment.vector_grad(f, argnums=0)(_x, _y)
939
- # pprint(g)
940
- # self.assertTrue(jnp.array_equal(g, 2 * _x))
941
- #
942
- # g = brainstate.augment.vector_grad(f, argnums=(0,))(_x, _y)
943
- # self.assertTrue(jnp.array_equal(g[0], 2 * _x))
944
- #
945
- # g = brainstate.augment.vector_grad(f, argnums=(0, 1))(_x, _y)
946
- # pprint(g)
947
- # self.assertTrue(jnp.array_equal(g[0], 2 * _x))
948
- # self.assertTrue(jnp.array_equal(g[1], 2 * _y))
949
- #
950
- # def test_aux1(self):
951
- # def f(x, y):
952
- # dx = x ** 2 + y ** 2 + 10
953
- # dy = x ** 3 + y ** 3 - 10
954
- # return dx, dy
955
- #
956
- # _x = jnp.ones(5)
957
- # _y = jnp.ones(5)
958
- #
959
- # g, aux = brainstate.augment.vector_grad(f, has_aux=True)(_x, _y)
960
- # pprint(g, )
961
- # pprint(aux)
962
- # self.assertTrue(jnp.array_equal(g, 2 * _x))
963
- # self.assertTrue(jnp.array_equal(aux, _x ** 3 + _y ** 3 - 10))
964
- #
965
- # def test_return1(self):
966
- # def f(x, y):
967
- # dx = x ** 2 + y ** 2 + 10
968
- # return dx
969
- #
970
- # _x = jnp.ones(5)
971
- # _y = jnp.ones(5)
972
- #
973
- # g, value = brainstate.augment.vector_grad(f, return_value=True)(_x, _y)
974
- # pprint(g, )
975
- # pprint(value)
976
- # self.assertTrue(jnp.array_equal(g, 2 * _x))
977
- # self.assertTrue(jnp.array_equal(value, _x ** 2 + _y ** 2 + 10))
978
- #
979
- # def test_return_aux1(self):
980
- # def f(x, y):
981
- # dx = x ** 2 + y ** 2 + 10
982
- # dy = x ** 3 + y ** 3 - 10
983
- # return dx, dy
984
- #
985
- # _x = jnp.ones(5)
986
- # _y = jnp.ones(5)
987
- #
988
- # g, value, aux = brainstate.augment.vector_grad(f, has_aux=True, return_value=True)(_x, _y)
989
- # print('grad', g)
990
- # print('value', value)
991
- # print('aux', aux)
992
- # self.assertTrue(jnp.array_equal(g, 2 * _x))
993
- # self.assertTrue(jnp.array_equal(value, _x ** 2 + _y ** 2 + 10))
994
- # self.assertTrue(jnp.array_equal(aux, _x ** 3 + _y ** 3 - 10))
995
- #
996
- #
997
- # class TestClassFuncVectorGrad(unittest.TestCase):
998
- # def test1(self):
999
- # class Test(brainstate.nn.Module):
1000
- # def __init__(self):
1001
- # super(Test, self).__init__()
1002
- # self.x = jnp.Variable(jnp.ones(5))
1003
- # self.y = jnp.Variable(jnp.ones(5))
1004
- #
1005
- # def __call__(self, *args, **kwargs):
1006
- # return self.x ** 2 + self.y ** 2 + 10
1007
- #
1008
- # t = Test()
1009
- #
1010
- # g = brainstate.augment.vector_grad(t, grad_states=t.x)()
1011
- # self.assertTrue(jnp.array_equal(g, 2 * t.x))
1012
- #
1013
- # g = brainstate.augment.vector_grad(t, grad_states=(t.x,))()
1014
- # self.assertTrue(jnp.array_equal(g[0], 2 * t.x))
1015
- #
1016
- # g = brainstate.augment.vector_grad(t, grad_states=(t.x, t.y))()
1017
- # self.assertTrue(jnp.array_equal(g[0], 2 * t.x))
1018
- # self.assertTrue(jnp.array_equal(g[1], 2 * t.y))
1019
- #
1020
- #
1021
- # def vgrad(f, *x):
1022
- # y, vjp_fn = jax.vjp(f, *x)
1023
- # return vjp_fn(jnp.ones(y.shape).value)[0]
1024
- #
1025
- #
1026
- # class TestDebug(parameterized.TestCase):
1027
- # def test_debug1(self):
1028
- # a = brainstate.random.RandomState()
1029
- #
1030
- # def f(b):
1031
- # print(a.value)
1032
- # return a + b + a.random()
1033
- #
1034
- # f = brainstate.augment.vector_grad(f, argnums=0)
1035
- # f(1.)
1036
- #
1037
- # with jax.disable_jit():
1038
- # f(1.)
1039
- #
1040
- # @parameterized.product(
1041
- # grad_fun=[brainstate.augment.grad, brainstate.augment.vector_grad]
1042
- # )
1043
- # def test_print_info1(self, grad_fun):
1044
- # file = tempfile.TemporaryFile(mode='w+')
1045
- #
1046
- # @functools.partial(grad_fun, argnums=0)
1047
- # def f2(a, b):
1048
- # print('compiling f2 ...', file=file)
1049
- # return a + b
1050
- #
1051
- # @functools.partial(grad_fun, argnums=0)
1052
- # def f1(a):
1053
- # print('compiling f1 ...', file=file)
1054
- # return f2(a, 1.)
1055
- #
1056
- # expect_res = '''
1057
- # compiling f1 ...
1058
- # compiling f2 ...
1059
- # compiling f1 ...
1060
- # compiling f2 ...
1061
- # '''
1062
- #
1063
- # print(f1(1.))
1064
- # file.seek(0)
1065
- # self.assertTrue(file.read().strip() == expect_res.strip())
1066
- #
1067
- # file = tempfile.TemporaryFile(mode='w+')
1068
- # with jax.disable_jit():
1069
- # expect_res = '''
1070
- # compiling f1 ...
1071
- # compiling f2 ...
1072
- # '''
1073
- # self.assertTrue(f1(1.) == 0.)
1074
- # file.seek(0)
1075
- # self.assertTrue(file.read().strip() == expect_res.strip())
1076
- #
1077
- # @parameterized.product(
1078
- # grad_fun=[brainstate.augment.grad, brainstate.augment.vector_grad]
1079
- # )
1080
- # def test_print_info2(self, grad_fun):
1081
- # file = tempfile.TemporaryFile(mode='w+')
1082
- #
1083
- # @functools.partial(grad_fun, argnums=0)
1084
- # def f1(a):
1085
- # @functools.partial(grad_fun, argnums=0)
1086
- # def f2(a, b):
1087
- # print('compiling f2 ...', file=file)
1088
- # return a + b
1089
- #
1090
- # print('compiling f1 ...', file=file)
1091
- # return f2(a, 1.)
1092
- #
1093
- # expect_res = '''
1094
- # compiling f1 ...
1095
- # compiling f2 ...
1096
- # compiling f1 ...
1097
- # compiling f2 ...
1098
- # compiling f2 ...
1099
- # '''
1100
- # self.assertTrue(f1(1.) == 0.)
1101
- # file.seek(0)
1102
- # self.assertTrue(file.read().strip() == expect_res.strip())
1103
- #
1104
- # file = tempfile.TemporaryFile(mode='w+')
1105
- # with jax.disable_jit():
1106
- # expect_res = '''
1107
- # compiling f1 ...
1108
- # compiling f2 ...
1109
- # '''
1110
- # self.assertTrue(f1(1.) == 0.)
1111
- # file.seek(0)
1112
- # # print(file.read().strip())
1113
- # self.assertTrue(file.read().strip() == expect_res.strip())
1114
- #
1115
- # def test_debug_correctness1(self):
1116
- # def test_f():
1117
- # a = jnp.Variable(jnp.ones(2))
1118
- # b = jnp.Variable(jnp.zeros(2))
1119
- #
1120
- # @brainstate.augment.vector_grad(argnums=0)
1121
- # def f1(c):
1122
- # a.value += 1
1123
- # b.value += 10
1124
- # return a * b * c
1125
- #
1126
- # return a, b, f1(1.)
1127
- #
1128
- # r1 = test_f()
1129
- # print(r1)
1130
- #
1131
- # with jax.disable_jit():
1132
- # r2 = test_f()
1133
- # print(r2)
1134
- # self.assertTrue(jnp.allclose(r1[0], r2[0]))
1135
- # self.assertTrue(jnp.allclose(r1[1], r2[1]))
1136
- # self.assertTrue(jnp.allclose(r1[2], r2[2]))
1137
- #
1138
- # def f1(c, a, b):
1139
- # a += 1
1140
- # b += 10
1141
- # return a * b * c
1142
- #
1143
- # r3 = vgrad(f1, 1., jnp.ones(2).value, jnp.zeros(2).value)
1144
- # self.assertTrue(jnp.allclose(r1[2], r3))
1145
- #
1146
- # def _bench_f2(self, dd):
1147
- # a = jnp.Variable(jnp.ones(2))
1148
- # b = jnp.Variable(jnp.zeros(2))
1149
- #
1150
- #
1151
- # def run_fun(d):
1152
- # @brainstate.augment.vector_grad(argnums=0)
1153
- # def f1(c):
1154
- # a.value += d
1155
- # b.value += 10
1156
- # return a * b * c
1157
- #
1158
- # return a, b, f1(1.)
1159
- #
1160
- # return run_fun(dd)
1161
- #
1162
- # def test_debug_correctness2(self):
1163
- # r1 = self._bench_f2(1.)
1164
- # print(r1)
1165
- #
1166
- # with jax.disable_jit():
1167
- # r2 = self._bench_f2(1.)
1168
- # print(r2)
1169
- #
1170
- # self.assertTrue(jnp.allclose(r1[0], r2[0]))
1171
- # self.assertTrue(jnp.allclose(r1[1], r2[1]))
1172
- # self.assertTrue(jnp.allclose(r1[2], r2[2]))
1173
- #
1174
- # def test_cache1(self):
1175
- # file = tempfile.TemporaryFile(mode='w+')
1176
- #
1177
- # def f(a, b):
1178
- # print('compiling f ...', file=file)
1179
- # return a + b
1180
- #
1181
- # grad1 = brainstate.augment.grad(f)(1., 2.) # call "f" twice, one for Variable finding, one for compiling
1182
- # grad2 = brainstate.augment.vector_grad(f)(1., 2.) # call "f" once for compiling
1183
- #
1184
- # file.seek(0)
1185
- # print(file.read().strip())
1186
- #
1187
- # expect_res = '''
1188
- # compiling f ...
1189
- # compiling f ...
1190
- # compiling f ...
1191
- # '''
1192
- # file.seek(0)
1193
- # self.assertTrue(file.read().strip() == expect_res.strip())
1194
- #
1195
- #
1196
-
1197
-
1198
- class TestUnitAwareGrad(unittest.TestCase):
1199
- def test_grad1(self):
1200
- def f(x):
1201
- return u.math.sum(x ** 2)
1202
-
1203
- x = jnp.array([1., 2., 3.]) * u.ms
1204
- g = brainstate.augment.grad(f, unit_aware=True)(x)
1205
- self.assertTrue(u.math.allclose(g, 2 * x))
1206
-
1207
- def test_vector_grad1(self):
1208
- def f(x):
1209
- return x ** 3
1210
-
1211
- x = jnp.array([1., 2., 3.]) * u.ms
1212
- g = brainstate.augment.vector_grad(f, unit_aware=True)(x)
1213
- self.assertTrue(u.math.allclose(g, 3 * x ** 2))
1214
-
1215
- def test_jacrev1(self):
1216
- def f(x, y):
1217
- return u.math.asarray([x[0] * y[0],
1218
- 5 * x[2] * y[1],
1219
- 4 * x[1] ** 2, ])
1220
-
1221
- _x = jnp.array([1., 2., 3.]) * u.ms
1222
- _y = jnp.array([10., 5.]) * u.ms
1223
-
1224
- g = brainstate.augment.jacrev(f, unit_aware=True, argnums=(0, 1))(_x, _y)
1225
- self.assertTrue(
1226
- u.math.allclose(
1227
- g[0],
1228
- u.math.asarray([
1229
- [10., 0., 0.],
1230
- [0., 0., 25.],
1231
- [0., 16., 0.]
1232
- ]) * u.ms
1233
- )
1234
- )
1235
-
1236
- self.assertTrue(
1237
- u.math.allclose(
1238
- g[1],
1239
- u.math.asarray([
1240
- [1., 0.],
1241
- [0., 15.],
1242
- [0., 0.]
1243
- ]) * u.ms
1244
- )
1245
- )
1246
-
1247
- def test_jacfwd1(self):
1248
- def f(x, y):
1249
- return u.math.asarray([x[0] * y[0],
1250
- 5 * x[2] * y[1],
1251
- 4 * x[1] ** 2, ])
1252
-
1253
- _x = jnp.array([1., 2., 3.]) * u.ms
1254
- _y = jnp.array([10., 5.]) * u.ms
1255
-
1256
- g = brainstate.augment.jacfwd(f, unit_aware=True, argnums=(0, 1))(_x, _y)
1257
- self.assertTrue(
1258
- u.math.allclose(
1259
- g[0],
1260
- u.math.asarray([
1261
- [10., 0., 0.],
1262
- [0., 0., 25.],
1263
- [0., 16., 0.]
1264
- ]) * u.ms
1265
- )
1266
- )
1267
-
1268
- self.assertTrue(
1269
- u.math.allclose(
1270
- g[1],
1271
- u.math.asarray([
1272
- [1., 0.],
1273
- [0., 15.],
1274
- [0., 0.]
1275
- ]) * u.ms
1276
- )
1277
- )
1278
-
1279
- def test_hessian(self):
1280
- unit = u.ms
1281
-
1282
- def scalar_function(x):
1283
- return x ** 3 + 3 * x * unit * unit + 2 * unit * unit * unit
1284
-
1285
- hess = brainstate.augment.hessian(scalar_function, unit_aware=True)
1286
- x = jnp.array(1.0) * unit
1287
- res = hess(x)
1288
- expected_hessian = jnp.array([[6.0]]) * unit
1289
- assert u.math.allclose(res, expected_hessian)
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
+ from pprint import pprint
20
+
21
+ import brainunit as u
22
+ import jax
23
+ import jax.numpy as jnp
24
+ import pytest
25
+
26
+ import brainstate
27
+ from brainstate.augment._autograd import _jacfwd
28
+
29
+
30
+ class TestPureFuncGrad(unittest.TestCase):
31
+ def test_grad_pure_func_1(self):
32
+ def call(a, b, c): return jnp.sum(a + b + c)
33
+
34
+ brainstate.random.seed(1)
35
+ a = jnp.ones(10)
36
+ b = brainstate.random.randn(10)
37
+ c = brainstate.random.uniform(size=10)
38
+ f_grad = brainstate.augment.grad(call, argnums=[0, 1, 2])
39
+ grads = f_grad(a, b, c)
40
+
41
+ for g in grads: assert (g == 1.).all()
42
+
43
+ def test_grad_pure_func_2(self):
44
+ def call(a, b, c): return jnp.sum(a + b + c)
45
+
46
+ brainstate.random.seed(1)
47
+ a = jnp.ones(10)
48
+ b = brainstate.random.randn(10)
49
+ c = brainstate.random.uniform(size=10)
50
+ f_grad = brainstate.augment.grad(call)
51
+ assert (f_grad(a, b, c) == 1.).all()
52
+
53
+ def test_grad_pure_func_aux1(self):
54
+ def call(a, b, c):
55
+ return jnp.sum(a + b + c), (jnp.sin(100), jnp.exp(0.1))
56
+
57
+ brainstate.random.seed(1)
58
+ f_grad = brainstate.augment.grad(call, argnums=[0, 1, 2])
59
+ with pytest.raises(TypeError):
60
+ f_grad(jnp.ones(10), brainstate.random.randn(10), brainstate.random.uniform(size=10))
61
+
62
+ def test_grad_pure_func_aux2(self):
63
+ def call(a, b, c):
64
+ return jnp.sum(a + b + c), (jnp.sin(100), jnp.exp(0.1))
65
+
66
+ brainstate.random.seed(1)
67
+ f_grad = brainstate.augment.grad(call, argnums=[0, 1, 2], has_aux=True)
68
+ grads, aux = f_grad(jnp.ones(10), brainstate.random.randn(10), brainstate.random.uniform(size=10))
69
+ for g in grads: assert (g == 1.).all()
70
+ assert aux[0] == jnp.sin(100)
71
+ assert aux[1] == jnp.exp(0.1)
72
+
73
+ def test_grad_pure_func_return1(self):
74
+ def call(a, b, c): return jnp.sum(a + b + c)
75
+
76
+ brainstate.random.seed(1)
77
+ a = jnp.ones(10)
78
+ b = brainstate.random.randn(10)
79
+ c = brainstate.random.uniform(size=10)
80
+ f_grad = brainstate.augment.grad(call, return_value=True)
81
+ grads, returns = f_grad(a, b, c)
82
+ assert (grads == 1.).all()
83
+ assert returns == jnp.sum(a + b + c)
84
+
85
+ def test_grad_func_return_aux1(self):
86
+ def call(a, b, c):
87
+ return jnp.sum(a + b + c), (jnp.sin(100), jnp.exp(0.1))
88
+
89
+ brainstate.random.seed(1)
90
+ a = jnp.ones(10)
91
+ b = brainstate.random.randn(10)
92
+ c = brainstate.random.uniform(size=10)
93
+ f_grad = brainstate.augment.grad(call, return_value=True, has_aux=True)
94
+ grads, returns, aux = f_grad(a, b, c)
95
+ assert (grads == 1.).all()
96
+ assert returns == jnp.sum(a + b + c)
97
+ assert aux[0] == jnp.sin(100)
98
+ assert aux[1] == jnp.exp(0.1)
99
+
100
+
101
+ class TestObjectFuncGrad(unittest.TestCase):
102
+ def test_grad_ob1(self):
103
+ class Test(brainstate.nn.Module):
104
+ def __init__(self):
105
+ super(Test, self).__init__()
106
+
107
+ self.a = brainstate.ParamState(jnp.ones(10))
108
+ self.b = brainstate.ParamState(brainstate.random.randn(10))
109
+ self.c = brainstate.ParamState(brainstate.random.uniform(size=10))
110
+
111
+ def __call__(self):
112
+ return jnp.sum(self.a.value + self.b.value + self.c.value)
113
+
114
+ brainstate.random.seed(0)
115
+
116
+ t = Test()
117
+ f_grad = brainstate.augment.grad(t, grad_states={'a': t.a, 'b': t.b, 'c': t.c})
118
+ grads = f_grad()
119
+ for g in grads.values():
120
+ assert (g == 1.).all()
121
+
122
+ t = Test()
123
+ f_grad = brainstate.augment.grad(t, grad_states=[t.a, t.b])
124
+ grads = f_grad()
125
+ for g in grads: assert (g == 1.).all()
126
+
127
+ t = Test()
128
+ f_grad = brainstate.augment.grad(t, grad_states=t.a)
129
+ grads = f_grad()
130
+ assert (grads == 1.).all()
131
+
132
+ t = Test()
133
+ f_grad = brainstate.augment.grad(t, grad_states=t.states())
134
+ grads = f_grad()
135
+ for g in grads.values():
136
+ assert (g == 1.).all()
137
+
138
+ def test_grad_ob_aux(self):
139
+ class Test(brainstate.nn.Module):
140
+ def __init__(self):
141
+ super(Test, self).__init__()
142
+ self.a = brainstate.ParamState(jnp.ones(10))
143
+ self.b = brainstate.ParamState(brainstate.random.randn(10))
144
+ self.c = brainstate.ParamState(brainstate.random.uniform(size=10))
145
+
146
+ def __call__(self):
147
+ return jnp.sum(self.a.value + self.b.value + self.c.value), (jnp.sin(100), jnp.exp(0.1))
148
+
149
+ brainstate.random.seed(0)
150
+ t = Test()
151
+ f_grad = brainstate.augment.grad(t, grad_states=[t.a, t.b], has_aux=True)
152
+ grads, aux = f_grad()
153
+ for g in grads: assert (g == 1.).all()
154
+ assert aux[0] == jnp.sin(100)
155
+ assert aux[1] == jnp.exp(0.1)
156
+
157
+ t = Test()
158
+ f_grad = brainstate.augment.grad(t, grad_states=t.a, has_aux=True)
159
+ grads, aux = f_grad()
160
+ assert (grads == 1.).all()
161
+ assert aux[0] == jnp.sin(100)
162
+ assert aux[1] == jnp.exp(0.1)
163
+
164
+ t = Test()
165
+ f_grad = brainstate.augment.grad(t, grad_states=t.states(), has_aux=True)
166
+ grads, aux = f_grad()
167
+ self.assertTrue(len(grads) == len(t.states()))
168
+
169
+ def test_grad_ob_return(self):
170
+ class Test(brainstate.nn.Module):
171
+ def __init__(self):
172
+ super(Test, self).__init__()
173
+ self.a = brainstate.ParamState(jnp.ones(10))
174
+ self.b = brainstate.ParamState(brainstate.random.randn(10))
175
+ self.c = brainstate.ParamState(brainstate.random.uniform(size=10))
176
+
177
+ def __call__(self):
178
+ return jnp.sum(self.a.value + self.b.value + self.c.value)
179
+
180
+ brainstate.random.seed(0)
181
+ t = Test()
182
+ f_grad = brainstate.augment.grad(t, grad_states=[t.a, t.b], return_value=True)
183
+ grads, returns = f_grad()
184
+ for g in grads: assert (g == 1.).all()
185
+ assert returns == t()
186
+
187
+ t = Test()
188
+ f_grad = brainstate.augment.grad(t, grad_states=t.a, return_value=True)
189
+ grads, returns = f_grad()
190
+ assert (grads == 1.).all()
191
+ assert returns == t()
192
+
193
+ def test_grad_ob_aux_return(self):
194
+ class Test(brainstate.nn.Module):
195
+ def __init__(self):
196
+ super(Test, self).__init__()
197
+ self.a = brainstate.ParamState(jnp.ones(10))
198
+ self.b = brainstate.ParamState(brainstate.random.randn(10))
199
+ self.c = brainstate.ParamState(brainstate.random.uniform(size=10))
200
+
201
+ def __call__(self):
202
+ return jnp.sum(self.a.value + self.b.value + self.c.value), (jnp.sin(100), jnp.exp(0.1))
203
+
204
+ brainstate.random.seed(0)
205
+ t = Test()
206
+ f_grad = brainstate.augment.grad(t, grad_states=[t.a, t.b], has_aux=True, return_value=True)
207
+ grads, returns, aux = f_grad()
208
+ for g in grads: assert (g == 1.).all()
209
+ assert returns == jnp.sum(t.a.value + t.b.value + t.c.value)
210
+ assert aux[0] == jnp.sin(100)
211
+ assert aux[1] == jnp.exp(0.1)
212
+
213
+ t = Test()
214
+ f_grad = brainstate.augment.grad(t, grad_states=t.a, has_aux=True, return_value=True)
215
+ grads, returns, aux = f_grad()
216
+ assert (grads == 1.).all()
217
+ assert returns == jnp.sum(t.a.value + t.b.value + t.c.value)
218
+ assert aux[0] == jnp.sin(100)
219
+ assert aux[1] == jnp.exp(0.1)
220
+
221
+ def test_grad_ob_argnums(self):
222
+ class Test(brainstate.nn.Module):
223
+ def __init__(self):
224
+ super(Test, self).__init__()
225
+ brainstate.random.seed()
226
+ self.a = brainstate.ParamState(jnp.ones(10))
227
+ self.b = brainstate.ParamState(brainstate.random.randn(10))
228
+ self.c = brainstate.ParamState(brainstate.random.uniform(size=10))
229
+
230
+ def __call__(self, d):
231
+ return jnp.sum(self.a.value + self.b.value + self.c.value + 2 * d)
232
+
233
+ brainstate.random.seed(0)
234
+
235
+ t = Test()
236
+ f_grad = brainstate.augment.grad(t, t.states(), argnums=0)
237
+ var_grads, arg_grads = f_grad(brainstate.random.random(10))
238
+ for g in var_grads.values(): assert (g == 1.).all()
239
+ assert (arg_grads == 2.).all()
240
+
241
+ t = Test()
242
+ f_grad = brainstate.augment.grad(t, t.states(), argnums=[0])
243
+ var_grads, arg_grads = f_grad(brainstate.random.random(10))
244
+ for g in var_grads.values(): assert (g == 1.).all()
245
+ assert (arg_grads[0] == 2.).all()
246
+
247
+ t = Test()
248
+ f_grad = brainstate.augment.grad(t, argnums=0)
249
+ arg_grads = f_grad(brainstate.random.random(10))
250
+ assert (arg_grads == 2.).all()
251
+
252
+ t = Test()
253
+ f_grad = brainstate.augment.grad(t, argnums=[0])
254
+ arg_grads = f_grad(brainstate.random.random(10))
255
+ assert (arg_grads[0] == 2.).all()
256
+
257
+ def test_grad_ob_argnums_aux(self):
258
+ class Test(brainstate.nn.Module):
259
+ def __init__(self):
260
+ super(Test, self).__init__()
261
+ self.a = brainstate.ParamState(jnp.ones(10))
262
+ self.b = brainstate.ParamState(brainstate.random.randn(10))
263
+ self.c = brainstate.ParamState(brainstate.random.uniform(size=10))
264
+
265
+ def __call__(self, d):
266
+ return jnp.sum(self.a.value + self.b.value + self.c.value + 2 * d), (jnp.sin(100), jnp.exp(0.1))
267
+
268
+ brainstate.random.seed(0)
269
+
270
+ t = Test()
271
+ f_grad = brainstate.augment.grad(t, grad_states=t.states(), argnums=0, has_aux=True)
272
+ (var_grads, arg_grads), aux = f_grad(brainstate.random.random(10))
273
+ for g in var_grads.values(): assert (g == 1.).all()
274
+ assert (arg_grads == 2.).all()
275
+ assert aux[0] == jnp.sin(100)
276
+ assert aux[1] == jnp.exp(0.1)
277
+
278
+ t = Test()
279
+ f_grad = brainstate.augment.grad(t, grad_states=t.states(), argnums=[0], has_aux=True)
280
+ (var_grads, arg_grads), aux = f_grad(brainstate.random.random(10))
281
+ for g in var_grads.values(): assert (g == 1.).all()
282
+ assert (arg_grads[0] == 2.).all()
283
+ assert aux[0] == jnp.sin(100)
284
+ assert aux[1] == jnp.exp(0.1)
285
+
286
+ t = Test()
287
+ f_grad = brainstate.augment.grad(t, argnums=0, has_aux=True)
288
+ arg_grads, aux = f_grad(brainstate.random.random(10))
289
+ assert (arg_grads == 2.).all()
290
+ assert aux[0] == jnp.sin(100)
291
+ assert aux[1] == jnp.exp(0.1)
292
+
293
+ t = Test()
294
+ f_grad = brainstate.augment.grad(t, argnums=[0], has_aux=True)
295
+ arg_grads, aux = f_grad(brainstate.random.random(10))
296
+ assert (arg_grads[0] == 2.).all()
297
+ assert aux[0] == jnp.sin(100)
298
+ assert aux[1] == jnp.exp(0.1)
299
+
300
+ def test_grad_ob_argnums_return(self):
301
+ class Test(brainstate.nn.Module):
302
+ def __init__(self):
303
+ super(Test, self).__init__()
304
+
305
+ self.a = brainstate.ParamState(jnp.ones(10))
306
+ self.b = brainstate.ParamState(brainstate.random.randn(10))
307
+ self.c = brainstate.ParamState(brainstate.random.uniform(size=10))
308
+
309
+ def __call__(self, d):
310
+ return jnp.sum(self.a.value + self.b.value + self.c.value + 2 * d)
311
+
312
+ brainstate.random.seed(0)
313
+
314
+ t = Test()
315
+ f_grad = brainstate.augment.grad(t, t.states(), argnums=0, return_value=True)
316
+ d = brainstate.random.random(10)
317
+ (var_grads, arg_grads), loss = f_grad(d)
318
+ for g in var_grads.values():
319
+ assert (g == 1.).all()
320
+ assert (arg_grads == 2.).all()
321
+ assert loss == t(d)
322
+
323
+ t = Test()
324
+ f_grad = brainstate.augment.grad(t, t.states(), argnums=[0], return_value=True)
325
+ d = brainstate.random.random(10)
326
+ (var_grads, arg_grads), loss = f_grad(d)
327
+ for g in var_grads.values():
328
+ assert (g == 1.).all()
329
+ assert (arg_grads[0] == 2.).all()
330
+ assert loss == t(d)
331
+
332
+ t = Test()
333
+ f_grad = brainstate.augment.grad(t, argnums=0, return_value=True)
334
+ d = brainstate.random.random(10)
335
+ arg_grads, loss = f_grad(d)
336
+ assert (arg_grads == 2.).all()
337
+ assert loss == t(d)
338
+
339
+ t = Test()
340
+ f_grad = brainstate.augment.grad(t, argnums=[0], return_value=True)
341
+ d = brainstate.random.random(10)
342
+ arg_grads, loss = f_grad(d)
343
+ assert (arg_grads[0] == 2.).all()
344
+ assert loss == t(d)
345
+
346
+ def test_grad_ob_argnums_aux_return(self):
347
+ class Test(brainstate.nn.Module):
348
+ def __init__(self):
349
+ super(Test, self).__init__()
350
+ self.a = brainstate.ParamState(jnp.ones(10))
351
+ self.b = brainstate.ParamState(brainstate.random.randn(10))
352
+ self.c = brainstate.ParamState(brainstate.random.uniform(size=10))
353
+
354
+ def __call__(self, d):
355
+ return jnp.sum(self.a.value + self.b.value + self.c.value + 2 * d), (jnp.sin(100), jnp.exp(0.1))
356
+
357
+ brainstate.random.seed(0)
358
+
359
+ t = Test()
360
+ f_grad = brainstate.augment.grad(t, grad_states=t.states(), argnums=0, has_aux=True, return_value=True)
361
+ d = brainstate.random.random(10)
362
+ (var_grads, arg_grads), loss, aux = f_grad(d)
363
+ for g in var_grads.values(): assert (g == 1.).all()
364
+ assert (arg_grads == 2.).all()
365
+ assert aux[0] == jnp.sin(100)
366
+ assert aux[1] == jnp.exp(0.1)
367
+ assert loss == t(d)[0]
368
+
369
+ t = Test()
370
+ f_grad = brainstate.augment.grad(t, grad_states=t.states(), argnums=[0], has_aux=True, return_value=True)
371
+ d = brainstate.random.random(10)
372
+ (var_grads, arg_grads), loss, aux = f_grad(d)
373
+ for g in var_grads.values(): assert (g == 1.).all()
374
+ assert (arg_grads[0] == 2.).all()
375
+ assert aux[0] == jnp.sin(100)
376
+ assert aux[1] == jnp.exp(0.1)
377
+ assert loss == t(d)[0]
378
+
379
+ t = Test()
380
+ f_grad = brainstate.augment.grad(t, argnums=0, has_aux=True, return_value=True)
381
+ d = brainstate.random.random(10)
382
+ arg_grads, loss, aux = f_grad(d)
383
+ assert (arg_grads == 2.).all()
384
+ assert aux[0] == jnp.sin(100)
385
+ assert aux[1] == jnp.exp(0.1)
386
+ assert loss == t(d)[0]
387
+
388
+ t = Test()
389
+ f_grad = brainstate.augment.grad(t, argnums=[0], has_aux=True, return_value=True)
390
+ d = brainstate.random.random(10)
391
+ arg_grads, loss, aux = f_grad(d)
392
+ assert (arg_grads[0] == 2.).all()
393
+ assert aux[0] == jnp.sin(100)
394
+ assert aux[1] == jnp.exp(0.1)
395
+ assert loss == t(d)[0]
396
+
397
+
398
+ class TestPureFuncJacobian(unittest.TestCase):
399
+ def test1(self):
400
+ jac, aux = _jacfwd(lambda x: (x ** 3, [x ** 2]), has_aux=True)(3.)
401
+ self.assertTrue(jax.numpy.allclose(jac, jax.jacfwd(lambda x: x ** 3)(3.)))
402
+ self.assertTrue(aux[0] == 9.)
403
+
404
+ def test_jacfwd_and_aux_nested(self):
405
+ def f(x):
406
+ jac, aux = _jacfwd(lambda x: (x ** 3, [x ** 3]), has_aux=True)(x)
407
+ return aux[0]
408
+
409
+ f2 = lambda x: x ** 3
410
+
411
+ self.assertEqual(_jacfwd(f)(4.), _jacfwd(f2)(4.))
412
+ self.assertEqual(jax.jit(_jacfwd(f))(4.), _jacfwd(f2)(4.))
413
+ self.assertEqual(jax.jit(_jacfwd(jax.jit(f)))(4.), _jacfwd(f2)(4.))
414
+
415
+ self.assertEqual(_jacfwd(f)(jnp.asarray(4.)), _jacfwd(f2)(jnp.asarray(4.)))
416
+ self.assertEqual(jax.jit(_jacfwd(f))(jnp.asarray(4.)), _jacfwd(f2)(jnp.asarray(4.)))
417
+ self.assertEqual(jax.jit(_jacfwd(jax.jit(f)))(jnp.asarray(4.)), _jacfwd(f2)(jnp.asarray(4.)))
418
+
419
+ def f(x):
420
+ jac, aux = _jacfwd(lambda x: (x ** 3, [x ** 3]), has_aux=True)(x)
421
+ return aux[0] * jnp.sin(x)
422
+
423
+ f2 = lambda x: x ** 3 * jnp.sin(x)
424
+
425
+ self.assertEqual(_jacfwd(f)(4.), _jacfwd(f2)(4.))
426
+ self.assertEqual(jax.jit(_jacfwd(f))(4.), _jacfwd(f2)(4.))
427
+ self.assertEqual(jax.jit(_jacfwd(jax.jit(f)))(4.), _jacfwd(f2)(4.))
428
+
429
+ self.assertEqual(_jacfwd(f)(jnp.asarray(4.)), _jacfwd(f2)(jnp.asarray(4.)))
430
+ self.assertEqual(jax.jit(_jacfwd(f))(jnp.asarray(4.)), _jacfwd(f2)(jnp.asarray(4.)))
431
+ self.assertEqual(jax.jit(_jacfwd(jax.jit(f)))(jnp.asarray(4.)), _jacfwd(f2)(jnp.asarray(4.)))
432
+
433
+ def test_jacrev1(self):
434
+ def f1(x, y):
435
+ r = jnp.asarray([x[0] * y[0], 5 * x[2] * y[1], 4 * x[1] ** 2 - 2 * x[2], x[2] * jnp.sin(x[0])])
436
+ return r
437
+
438
+ br = brainstate.augment.jacrev(f1)(jnp.array([1., 2., 3.]), jnp.array([10., 5.]))
439
+ jr = jax.jacrev(f1)(jnp.array([1., 2., 3.]), jnp.array([10., 5.]))
440
+ assert (br == jr).all()
441
+
442
+ br = brainstate.augment.jacrev(f1, argnums=(0, 1))(jnp.array([1., 2., 3.]), jnp.array([10., 5.]))
443
+ jr = jax.jacrev(f1, argnums=(0, 1))(jnp.array([1., 2., 3.]), jnp.array([10., 5.]))
444
+ assert (br[0] == jr[0]).all()
445
+ assert (br[1] == jr[1]).all()
446
+
447
+ def test_jacrev2(self):
448
+ print()
449
+
450
+ def f2(x, y):
451
+ r1 = jnp.asarray([x[0] * y[0], 5 * x[2] * y[1]])
452
+ r2 = jnp.asarray([4 * x[1] ** 2 - 2 * x[2], x[2] * jnp.sin(x[0])])
453
+ return r1, r2
454
+
455
+ jr = jax.jacrev(f2)(jnp.array([1., 2., 3.]), jnp.array([10., 5.]))
456
+ pprint(jr)
457
+
458
+ br = brainstate.augment.jacrev(f2)(jnp.array([1., 2., 3.]), jnp.array([10., 5.]))
459
+ pprint(br)
460
+ assert jnp.array_equal(br[0], jr[0])
461
+ assert jnp.array_equal(br[1], jr[1])
462
+
463
+ br = brainstate.augment.jacrev(f2)(jnp.array([1., 2., 3.]), jnp.array([10., 5.]))
464
+ pprint(br)
465
+ assert jnp.array_equal(br[0], jr[0])
466
+ assert jnp.array_equal(br[1], jr[1])
467
+
468
+ def f2(x, y):
469
+ r1 = jnp.asarray([x[0] * y[0], 5 * x[2] * y[1]])
470
+ r2 = jnp.asarray([4 * x[1] ** 2 - 2 * x[2], x[2] * jnp.sin(x[0])])
471
+ return r1, r2
472
+
473
+ br = brainstate.augment.jacrev(f2)(jnp.array([1., 2., 3.]), jnp.array([10., 5.]))
474
+ pprint(br)
475
+ assert jnp.array_equal(br[0], jr[0])
476
+ assert jnp.array_equal(br[1], jr[1])
477
+
478
+ br = brainstate.augment.jacrev(f2)(jnp.array([1., 2., 3.]), jnp.array([10., 5.]))
479
+ pprint(br)
480
+ assert jnp.array_equal(br[0], jr[0])
481
+ assert jnp.array_equal(br[1], jr[1])
482
+
483
+ def test_jacrev3(self):
484
+ print()
485
+
486
+ def f3(x, y):
487
+ r1 = jnp.asarray([x[0] * y[0], 5 * x[2] * y[1]])
488
+ r2 = jnp.asarray([4 * x[1] ** 2 - 2 * x[2], x[2] * jnp.sin(x[0])])
489
+ return r1, r2
490
+
491
+ jr = jax.jacrev(f3, argnums=(0, 1))(jnp.array([1., 2., 3.]), jnp.array([10., 5.]))
492
+ pprint(jr)
493
+
494
+ br = brainstate.augment.jacrev(f3, argnums=(0, 1))(jnp.array([1., 2., 3.]), jnp.array([10., 5.]))
495
+ pprint(br)
496
+ assert jnp.array_equal(br[0][0], jr[0][0])
497
+ assert jnp.array_equal(br[0][1], jr[0][1])
498
+ assert jnp.array_equal(br[1][0], jr[1][0])
499
+ assert jnp.array_equal(br[1][1], jr[1][1])
500
+
501
+ br = brainstate.augment.jacrev(f3, argnums=(0, 1))(jnp.array([1., 2., 3.]), jnp.array([10., 5.]))
502
+ pprint(br)
503
+ assert jnp.array_equal(br[0][0], jr[0][0])
504
+ assert jnp.array_equal(br[0][1], jr[0][1])
505
+ assert jnp.array_equal(br[1][0], jr[1][0])
506
+ assert jnp.array_equal(br[1][1], jr[1][1])
507
+
508
+ def f3(x, y):
509
+ r1 = jnp.asarray([x[0] * y[0], 5 * x[2] * y[1]])
510
+ r2 = jnp.asarray([4 * x[1] ** 2 - 2 * x[2], x[2] * jnp.sin(x[0])])
511
+ return r1, r2
512
+
513
+ br = brainstate.augment.jacrev(f3, argnums=(0, 1))(jnp.array([1., 2., 3.]), jnp.array([10., 5.]))
514
+ pprint(br)
515
+ assert jnp.array_equal(br[0][0], jr[0][0])
516
+ assert jnp.array_equal(br[0][1], jr[0][1])
517
+ assert jnp.array_equal(br[1][0], jr[1][0])
518
+ assert jnp.array_equal(br[1][1], jr[1][1])
519
+
520
+ br = brainstate.augment.jacrev(f3, argnums=(0, 1))(jnp.array([1., 2., 3.]), jnp.array([10., 5.]))
521
+ pprint(br)
522
+ assert jnp.array_equal(br[0][0], jr[0][0])
523
+ assert jnp.array_equal(br[0][1], jr[0][1])
524
+ assert jnp.array_equal(br[1][0], jr[1][0])
525
+ assert jnp.array_equal(br[1][1], jr[1][1])
526
+
527
+ def test_jacrev_aux1(self):
528
+ x = jnp.array([1., 2., 3.])
529
+ y = jnp.array([10., 5.])
530
+
531
+ def f1(x, y):
532
+ a = 4 * x[1] ** 2 - 2 * x[2]
533
+ r = jnp.asarray([x[0] * y[0], 5 * x[2] * y[1], a, x[2] * jnp.sin(x[0])])
534
+ return r, a
535
+
536
+ f2 = lambda *args: f1(*args)[0]
537
+ jr = jax.jacrev(f2)(x, y) # jax jacobian
538
+ pprint(jr)
539
+ grads, aux = brainstate.augment.jacrev(f1, has_aux=True)(x, y)
540
+ assert (grads == jr).all()
541
+ assert aux == (4 * x[1] ** 2 - 2 * x[2])
542
+
543
+ jr = jax.jacrev(f2, argnums=(0, 1))(x, y) # jax jacobian
544
+ pprint(jr)
545
+ grads, aux = brainstate.augment.jacrev(f1, argnums=(0, 1), has_aux=True)(x, y)
546
+ assert (grads[0] == jr[0]).all()
547
+ assert (grads[1] == jr[1]).all()
548
+ assert aux == (4 * x[1] ** 2 - 2 * x[2])
549
+
550
+ def test_jacrev_return_aux1(self):
551
+ with brainstate.environ.context(precision=64):
552
+ def f1(x, y):
553
+ a = 4 * x[1] ** 2 - 2 * x[2]
554
+ r = jnp.asarray([x[0] * y[0], 5 * x[2] * y[1], a, x[2] * jnp.sin(x[0])])
555
+ return r, a
556
+
557
+ _x = jnp.array([1., 2., 3.])
558
+ _y = jnp.array([10., 5.])
559
+ _r, _a = f1(_x, _y)
560
+ f2 = lambda *args: f1(*args)[0]
561
+ _g1 = jax.jacrev(f2)(_x, _y) # jax jacobian
562
+ pprint(_g1)
563
+ _g2 = jax.jacrev(f2, argnums=(0, 1))(_x, _y) # jax jacobian
564
+ pprint(_g2)
565
+
566
+ grads, vec, aux = brainstate.augment.jacrev(f1, return_value=True, has_aux=True)(_x, _y)
567
+ assert (grads == _g1).all()
568
+ assert aux == _a
569
+ assert (vec == _r).all()
570
+
571
+ grads, vec, aux = brainstate.augment.jacrev(f1, return_value=True, argnums=(0, 1), has_aux=True)(_x, _y)
572
+ assert (grads[0] == _g2[0]).all()
573
+ assert (grads[1] == _g2[1]).all()
574
+ assert aux == _a
575
+ assert (vec == _r).all()
576
+
577
+
578
+ class TestClassFuncJacobian(unittest.TestCase):
579
+ def test_jacrev1(self):
580
+ def f1(x, y):
581
+ r = jnp.asarray([x[0] * y[0], 5 * x[2] * y[1], 4 * x[1] ** 2 - 2 * x[2], x[2] * jnp.sin(x[0])])
582
+ return r
583
+
584
+ _x = jnp.array([1., 2., 3.])
585
+ _y = jnp.array([10., 5.])
586
+
587
+ class Test(brainstate.nn.Module):
588
+ def __init__(self):
589
+ super(Test, self).__init__()
590
+ self.x = brainstate.State(jnp.array([1., 2., 3.]))
591
+ self.y = brainstate.State(jnp.array([10., 5.]))
592
+
593
+ def __call__(self, ):
594
+ a = self.x.value[0] * self.y.value[0]
595
+ b = 5 * self.x.value[2] * self.y.value[1]
596
+ c = 4 * self.x.value[1] ** 2 - 2 * self.x.value[2]
597
+ d = self.x.value[2] * jnp.sin(self.x.value[0])
598
+ r = jnp.asarray([a, b, c, d])
599
+ return r
600
+
601
+ _jr = jax.jacrev(f1)(_x, _y)
602
+ t = Test()
603
+ br = brainstate.augment.jacrev(t, grad_states=t.x)()
604
+ self.assertTrue((br == _jr).all())
605
+
606
+ _jr = jax.jacrev(f1, argnums=(0, 1))(_x, _y)
607
+ t = Test()
608
+ br = brainstate.augment.jacrev(t, grad_states=[t.x, t.y])()
609
+ self.assertTrue((br[0] == _jr[0]).all())
610
+ self.assertTrue((br[1] == _jr[1]).all())
611
+
612
+
613
+ #
614
+ # def test_jacfwd1(self):
615
+ # def f1(x, y):
616
+ # r = jnp.asarray([x[0] * y[0], 5 * x[2] * y[1], 4 * x[1] ** 2 - 2 * x[2], x[2] * jnp.sin(x[0])])
617
+ # return r
618
+ #
619
+ # _x = jnp.array([1., 2., 3.])
620
+ # _y = jnp.array([10., 5.])
621
+ #
622
+ # class Test(brainstate.nn.Module):
623
+ # def __init__(self):
624
+ # super(Test, self).__init__()
625
+ # self.x = jnp.Variable(jnp.array([1., 2., 3.]))
626
+ # self.y = jnp.Variable(jnp.array([10., 5.]))
627
+ #
628
+ # def __call__(self, ):
629
+ # a = self.x[0] * self.y[0]
630
+ # b = 5 * self.x[2] * self.y[1]
631
+ # c = 4 * self.x[1] ** 2 - 2 * self.x[2]
632
+ # d = self.x[2] * jnp.sin(self.x[0])
633
+ # r = jnp.asarray([a, b, c, d])
634
+ # return r
635
+ #
636
+ # _jr = jax.jacfwd(f1)(_x, _y)
637
+ # t = Test()
638
+ # br = brainstate.augment.jacfwd(t, grad_states=t.x)()
639
+ # self.assertTrue((br == _jr).all())
640
+ #
641
+ # _jr = jax.jacfwd(f1, argnums=(0, 1))(_x, _y)
642
+ # t = Test()
643
+ # br = brainstate.augment.jacfwd(t, grad_states=[t.x, t.y])()
644
+ # self.assertTrue((br[0] == _jr[0]).all())
645
+ # self.assertTrue((br[1] == _jr[1]).all())
646
+ #
647
+ # def test_jacrev2(self):
648
+ # def f1(x, y):
649
+ # r = jnp.asarray([x[0] * y[0], 5 * x[2] * y[1], 4 * x[1] ** 2 - 2 * x[2], x[2] * jnp.sin(x[0])])
650
+ # return r
651
+ #
652
+ # _x = jnp.array([1., 2., 3.])
653
+ # _y = jnp.array([10., 5.])
654
+ #
655
+ # class Test(brainstate.nn.Module):
656
+ # def __init__(self):
657
+ # super(Test, self).__init__()
658
+ # self.x = jnp.Variable(jnp.array([1., 2., 3.]))
659
+ #
660
+ # def __call__(self, y):
661
+ # a = self.x[0] * y[0]
662
+ # b = 5 * self.x[2] * y[1]
663
+ # c = 4 * self.x[1] ** 2 - 2 * self.x[2]
664
+ # d = self.x[2] * jnp.sin(self.x[0])
665
+ # r = jnp.asarray([a, b, c, d])
666
+ # return r
667
+ #
668
+ # _jr = jax.jacrev(f1)(_x, _y)
669
+ # t = Test()
670
+ # br = brainstate.augment.jacrev(t, grad_states=t.x)(_y)
671
+ # self.assertTrue((br == _jr).all())
672
+ #
673
+ # _jr = jax.jacrev(f1, argnums=(0, 1))(_x, _y)
674
+ # t = Test()
675
+ # var_grads, arg_grads = brainstate.augment.jacrev(t, grad_states=t.x, argnums=0)(_y)
676
+ # print(var_grads, )
677
+ # print(arg_grads, )
678
+ # self.assertTrue((var_grads == _jr[0]).all())
679
+ # self.assertTrue((arg_grads == _jr[1]).all())
680
+ #
681
+ # def test_jacfwd2(self):
682
+ # def f1(x, y):
683
+ # r = jnp.asarray([x[0] * y[0], 5 * x[2] * y[1], 4 * x[1] ** 2 - 2 * x[2], x[2] * jnp.sin(x[0])])
684
+ # return r
685
+ #
686
+ # _x = jnp.array([1., 2., 3.])
687
+ # _y = jnp.array([10., 5.])
688
+ #
689
+ # class Test(brainstate.nn.Module):
690
+ # def __init__(self):
691
+ # super(Test, self).__init__()
692
+ # self.x = jnp.Variable(jnp.array([1., 2., 3.]))
693
+ #
694
+ # def __call__(self, y):
695
+ # a = self.x[0] * y[0]
696
+ # b = 5 * self.x[2] * y[1]
697
+ # c = 4 * self.x[1] ** 2 - 2 * self.x[2]
698
+ # d = self.x[2] * jnp.sin(self.x[0])
699
+ # r = jnp.asarray([a, b, c, d])
700
+ # return r
701
+ #
702
+ # _jr = jax.jacfwd(f1)(_x, _y)
703
+ # t = Test()
704
+ # br = brainstate.augment.jacfwd(t, grad_states=t.x)(_y)
705
+ # self.assertTrue((br == _jr).all())
706
+ #
707
+ # _jr = jax.jacfwd(f1, argnums=(0, 1))(_x, _y)
708
+ # t = Test()
709
+ # var_grads, arg_grads = brainstate.augment.jacfwd(t, grad_states=t.x, argnums=0)(_y)
710
+ # print(var_grads, )
711
+ # print(arg_grads, )
712
+ # self.assertTrue((var_grads == _jr[0]).all())
713
+ # self.assertTrue((arg_grads == _jr[1]).all())
714
+ #
715
+ # def test_jacrev_aux1(self):
716
+ # jnp.enable_x64()
717
+ #
718
+ # def f1(x, y):
719
+ # r = jnp.asarray([x[0] * y[0], 5 * x[2] * y[1], 4 * x[1] ** 2 - 2 * x[2], x[2] * jnp.sin(x[0])])
720
+ # return r
721
+ #
722
+ # _x = jnp.array([1., 2., 3.])
723
+ # _y = jnp.array([10., 5.])
724
+ #
725
+ # class Test(brainstate.nn.Module):
726
+ # def __init__(self):
727
+ # super(Test, self).__init__()
728
+ # self.x = jnp.Variable(jnp.array([1., 2., 3.]))
729
+ #
730
+ # def __call__(self, y):
731
+ # a = self.x[0] * y[0]
732
+ # b = 5 * self.x[2] * y[1]
733
+ # c = 4 * self.x[1] ** 2 - 2 * self.x[2]
734
+ # d = self.x[2] * jnp.sin(self.x[0])
735
+ # r = jnp.asarray([a, b, c, d])
736
+ # return r, (c, d)
737
+ #
738
+ # _jr = jax.jacrev(f1)(_x, _y)
739
+ # t = Test()
740
+ # br, _ = brainstate.augment.jacrev(t, grad_states=t.x, has_aux=True)(_y)
741
+ # self.assertTrue((br == _jr).all())
742
+ #
743
+ # t = Test()
744
+ # _jr = jax.jacrev(f1, argnums=(0, 1))(_x, _y)
745
+ # _aux = t(_y)[1]
746
+ # (var_grads, arg_grads), aux = brainstate.augment.jacrev(t, grad_states=t.x, argnums=0, has_aux=True)(_y)
747
+ # print(var_grads, )
748
+ # print(arg_grads, )
749
+ # self.assertTrue((var_grads == _jr[0]).all())
750
+ # self.assertTrue((arg_grads == _jr[1]).all())
751
+ # self.assertTrue(jnp.array_equal(aux, _aux))
752
+ #
753
+ # jnp.disable_x64()
754
+ #
755
+ # def test_jacfwd_aux1(self):
756
+ # jnp.enable_x64()
757
+ #
758
+ # def f1(x, y):
759
+ # r = jnp.asarray([x[0] * y[0], 5 * x[2] * y[1], 4 * x[1] ** 2 - 2 * x[2], x[2] * jnp.sin(x[0])])
760
+ # return r
761
+ #
762
+ # _x = jnp.array([1., 2., 3.])
763
+ # _y = jnp.array([10., 5.])
764
+ #
765
+ # class Test(brainstate.nn.Module):
766
+ # def __init__(self):
767
+ # super(Test, self).__init__()
768
+ # self.x = jnp.Variable(jnp.array([1., 2., 3.]))
769
+ #
770
+ # def __call__(self, y):
771
+ # a = self.x[0] * y[0]
772
+ # b = 5 * self.x[2] * y[1]
773
+ # c = 4 * self.x[1] ** 2 - 2 * self.x[2]
774
+ # d = self.x[2] * jnp.sin(self.x[0])
775
+ # r = jnp.asarray([a, b, c, d])
776
+ # return r, (c, d)
777
+ #
778
+ # _jr = jax.jacfwd(f1)(_x, _y)
779
+ # t = Test()
780
+ # br, (c, d) = brainstate.augment.jacfwd(t, grad_states=t.x, has_aux=True)(_y)
781
+ # # print(_jr)
782
+ # # print(br)
783
+ # a = (br == _jr)
784
+ # self.assertTrue(a.all())
785
+ #
786
+ # t = Test()
787
+ # _jr = jax.jacfwd(f1, argnums=(0, 1))(_x, _y)
788
+ # _aux = t(_y)[1]
789
+ # (var_grads, arg_grads), aux = brainstate.augment.jacfwd(t, grad_states=t.x, argnums=0, has_aux=True)(_y)
790
+ # print(var_grads, )
791
+ # print(arg_grads, )
792
+ # self.assertTrue((var_grads == _jr[0]).all())
793
+ # self.assertTrue((arg_grads == _jr[1]).all())
794
+ # self.assertTrue(jnp.array_equal(aux, _aux))
795
+ #
796
+ # jnp.disable_x64()
797
+ #
798
+ # def test_jacrev_return_aux1(self):
799
+ # jnp.enable_x64()
800
+ #
801
+ # def f1(x, y):
802
+ # r = jnp.asarray([x[0] * y[0], 5 * x[2] * y[1], 4 * x[1] ** 2 - 2 * x[2], x[2] * jnp.sin(x[0])])
803
+ # return r
804
+ #
805
+ # _x = jnp.array([1., 2., 3.])
806
+ # _y = jnp.array([10., 5.])
807
+ #
808
+ # class Test(brainstate.nn.Module):
809
+ # def __init__(self):
810
+ # super(Test, self).__init__()
811
+ # self.x = jnp.Variable(jnp.array([1., 2., 3.]))
812
+ #
813
+ # def __call__(self, y):
814
+ # a = self.x[0] * y[0]
815
+ # b = 5 * self.x[2] * y[1]
816
+ # c = 4 * self.x[1] ** 2 - 2 * self.x[2]
817
+ # d = self.x[2] * jnp.sin(self.x[0])
818
+ # r = jnp.asarray([a, b, c, d])
819
+ # return r, (c, d)
820
+ #
821
+ # _jr = jax.jacrev(f1)(_x, _y)
822
+ # t = Test()
823
+ # br, _ = brainstate.augment.jacrev(t, grad_states=t.x, has_aux=True)(_y)
824
+ # self.assertTrue((br == _jr).all())
825
+ #
826
+ # t = Test()
827
+ # _jr = jax.jacrev(f1, argnums=(0, 1))(_x, _y)
828
+ # _val, _aux = t(_y)
829
+ # (var_grads, arg_grads), value, aux = brainstate.augment.jacrev(t, grad_states=t.x, argnums=0, has_aux=True, return_value=True)(_y)
830
+ # print(var_grads, )
831
+ # print(arg_grads, )
832
+ # self.assertTrue((var_grads == _jr[0]).all())
833
+ # self.assertTrue((arg_grads == _jr[1]).all())
834
+ # self.assertTrue(jnp.array_equal(aux, _aux))
835
+ # self.assertTrue(jnp.array_equal(value, _val))
836
+ #
837
+ # jnp.disable_x64()
838
+ #
839
+ # def test_jacfwd_return_aux1(self):
840
+ # jnp.enable_x64()
841
+ #
842
+ # def f1(x, y):
843
+ # r = jnp.asarray([x[0] * y[0], 5 * x[2] * y[1], 4 * x[1] ** 2 - 2 * x[2], x[2] * jnp.sin(x[0])])
844
+ # return r
845
+ #
846
+ # _x = jnp.array([1., 2., 3.])
847
+ # _y = jnp.array([10., 5.])
848
+ #
849
+ # class Test(brainstate.nn.Module):
850
+ # def __init__(self):
851
+ # super(Test, self).__init__()
852
+ # self.x = jnp.Variable(jnp.array([1., 2., 3.]))
853
+ #
854
+ # def __call__(self, y):
855
+ # a = self.x[0] * y[0]
856
+ # b = 5 * self.x[2] * y[1]
857
+ # c = 4 * self.x[1] ** 2 - 2 * self.x[2]
858
+ # d = self.x[2] * jnp.sin(self.x[0])
859
+ # r = jnp.asarray([a, b, c, d])
860
+ # return r, (c, d)
861
+ #
862
+ # _jr = jax.jacfwd(f1)(_x, _y)
863
+ # t = Test()
864
+ # br, _ = brainstate.augment.jacfwd(t, grad_states=t.x, has_aux=True)(_y)
865
+ # self.assertTrue((br == _jr).all())
866
+ #
867
+ # t = Test()
868
+ # _jr = jax.jacfwd(f1, argnums=(0, 1))(_x, _y)
869
+ # _val, _aux = t(_y)
870
+ # (var_grads, arg_grads), value, aux = brainstate.augment.jacfwd(t, grad_states=t.x, argnums=0, has_aux=True, return_value=True)(_y)
871
+ # print(_val, )
872
+ # print('_aux: ', _aux, 'aux: ', aux)
873
+ # print(var_grads, )
874
+ # print(arg_grads, )
875
+ # self.assertTrue((var_grads == _jr[0]).all())
876
+ # self.assertTrue((arg_grads == _jr[1]).all())
877
+ # self.assertTrue(jnp.array_equal(aux, _aux))
878
+ # self.assertTrue(jnp.array_equal(value, _val))
879
+ #
880
+ # jnp.disable_x64()
881
+ #
882
+ #
883
+ # class TestPureFuncVectorGrad(unittest.TestCase):
884
+ # def test1(self):
885
+ # f = lambda x: 3 * x ** 2
886
+ # _x = jnp.ones(10)
887
+ # pprint(brainstate.augment.vector_grad(f, argnums=0)(_x))
888
+ #
889
+ # def test2(self):
890
+ # def f(x, y):
891
+ # dx = x ** 2 + y ** 2 + 10
892
+ # return dx
893
+ #
894
+ # _x = jnp.ones(5)
895
+ # _y = jnp.ones(5)
896
+ #
897
+ # g = brainstate.augment.vector_grad(f, argnums=0)(_x, _y)
898
+ # pprint(g)
899
+ # self.assertTrue(jnp.array_equal(g, 2 * _x))
900
+ #
901
+ # g = brainstate.augment.vector_grad(f, argnums=(0,))(_x, _y)
902
+ # self.assertTrue(jnp.array_equal(g[0], 2 * _x))
903
+ #
904
+ # g = brainstate.augment.vector_grad(f, argnums=(0, 1))(_x, _y)
905
+ # pprint(g)
906
+ # self.assertTrue(jnp.array_equal(g[0], 2 * _x))
907
+ # self.assertTrue(jnp.array_equal(g[1], 2 * _y))
908
+ #
909
+ # def test3(self):
910
+ # def f(x, y):
911
+ # dx = x ** 2 + y ** 2 + 10
912
+ # dy = x ** 3 + y ** 3 - 10
913
+ # return dx, dy
914
+ #
915
+ # _x = jnp.ones(5)
916
+ # _y = jnp.ones(5)
917
+ #
918
+ # g = brainstate.augment.vector_grad(f, argnums=0)(_x, _y)
919
+ # # pprint(g)
920
+ # self.assertTrue(jnp.array_equal(g, 2 * _x + 3 * _x ** 2))
921
+ #
922
+ # g = brainstate.augment.vector_grad(f, argnums=(0,))(_x, _y)
923
+ # self.assertTrue(jnp.array_equal(g[0], 2 * _x + 3 * _x ** 2))
924
+ #
925
+ # g = brainstate.augment.vector_grad(f, argnums=(0, 1))(_x, _y)
926
+ # # pprint(g)
927
+ # self.assertTrue(jnp.array_equal(g[0], 2 * _x + 3 * _x ** 2))
928
+ # self.assertTrue(jnp.array_equal(g[1], 2 * _y + 3 * _y ** 2))
929
+ #
930
+ # def test4_2d(self):
931
+ # def f(x, y):
932
+ # dx = x ** 2 + y ** 2 + 10
933
+ # return dx
934
+ #
935
+ # _x = jnp.ones((5, 5))
936
+ # _y = jnp.ones((5, 5))
937
+ #
938
+ # g = brainstate.augment.vector_grad(f, argnums=0)(_x, _y)
939
+ # pprint(g)
940
+ # self.assertTrue(jnp.array_equal(g, 2 * _x))
941
+ #
942
+ # g = brainstate.augment.vector_grad(f, argnums=(0,))(_x, _y)
943
+ # self.assertTrue(jnp.array_equal(g[0], 2 * _x))
944
+ #
945
+ # g = brainstate.augment.vector_grad(f, argnums=(0, 1))(_x, _y)
946
+ # pprint(g)
947
+ # self.assertTrue(jnp.array_equal(g[0], 2 * _x))
948
+ # self.assertTrue(jnp.array_equal(g[1], 2 * _y))
949
+ #
950
+ # def test_aux1(self):
951
+ # def f(x, y):
952
+ # dx = x ** 2 + y ** 2 + 10
953
+ # dy = x ** 3 + y ** 3 - 10
954
+ # return dx, dy
955
+ #
956
+ # _x = jnp.ones(5)
957
+ # _y = jnp.ones(5)
958
+ #
959
+ # g, aux = brainstate.augment.vector_grad(f, has_aux=True)(_x, _y)
960
+ # pprint(g, )
961
+ # pprint(aux)
962
+ # self.assertTrue(jnp.array_equal(g, 2 * _x))
963
+ # self.assertTrue(jnp.array_equal(aux, _x ** 3 + _y ** 3 - 10))
964
+ #
965
+ # def test_return1(self):
966
+ # def f(x, y):
967
+ # dx = x ** 2 + y ** 2 + 10
968
+ # return dx
969
+ #
970
+ # _x = jnp.ones(5)
971
+ # _y = jnp.ones(5)
972
+ #
973
+ # g, value = brainstate.augment.vector_grad(f, return_value=True)(_x, _y)
974
+ # pprint(g, )
975
+ # pprint(value)
976
+ # self.assertTrue(jnp.array_equal(g, 2 * _x))
977
+ # self.assertTrue(jnp.array_equal(value, _x ** 2 + _y ** 2 + 10))
978
+ #
979
+ # def test_return_aux1(self):
980
+ # def f(x, y):
981
+ # dx = x ** 2 + y ** 2 + 10
982
+ # dy = x ** 3 + y ** 3 - 10
983
+ # return dx, dy
984
+ #
985
+ # _x = jnp.ones(5)
986
+ # _y = jnp.ones(5)
987
+ #
988
+ # g, value, aux = brainstate.augment.vector_grad(f, has_aux=True, return_value=True)(_x, _y)
989
+ # print('grad', g)
990
+ # print('value', value)
991
+ # print('aux', aux)
992
+ # self.assertTrue(jnp.array_equal(g, 2 * _x))
993
+ # self.assertTrue(jnp.array_equal(value, _x ** 2 + _y ** 2 + 10))
994
+ # self.assertTrue(jnp.array_equal(aux, _x ** 3 + _y ** 3 - 10))
995
+ #
996
+ #
997
+ # class TestClassFuncVectorGrad(unittest.TestCase):
998
+ # def test1(self):
999
+ # class Test(brainstate.nn.Module):
1000
+ # def __init__(self):
1001
+ # super(Test, self).__init__()
1002
+ # self.x = jnp.Variable(jnp.ones(5))
1003
+ # self.y = jnp.Variable(jnp.ones(5))
1004
+ #
1005
+ # def __call__(self, *args, **kwargs):
1006
+ # return self.x ** 2 + self.y ** 2 + 10
1007
+ #
1008
+ # t = Test()
1009
+ #
1010
+ # g = brainstate.augment.vector_grad(t, grad_states=t.x)()
1011
+ # self.assertTrue(jnp.array_equal(g, 2 * t.x))
1012
+ #
1013
+ # g = brainstate.augment.vector_grad(t, grad_states=(t.x,))()
1014
+ # self.assertTrue(jnp.array_equal(g[0], 2 * t.x))
1015
+ #
1016
+ # g = brainstate.augment.vector_grad(t, grad_states=(t.x, t.y))()
1017
+ # self.assertTrue(jnp.array_equal(g[0], 2 * t.x))
1018
+ # self.assertTrue(jnp.array_equal(g[1], 2 * t.y))
1019
+ #
1020
+ #
1021
+ # def vgrad(f, *x):
1022
+ # y, vjp_fn = jax.vjp(f, *x)
1023
+ # return vjp_fn(jnp.ones(y.shape).value)[0]
1024
+ #
1025
+ #
1026
+ # class TestDebug(parameterized.TestCase):
1027
+ # def test_debug1(self):
1028
+ # a = brainstate.random.RandomState()
1029
+ #
1030
+ # def f(b):
1031
+ # print(a.value)
1032
+ # return a + b + a.random()
1033
+ #
1034
+ # f = brainstate.augment.vector_grad(f, argnums=0)
1035
+ # f(1.)
1036
+ #
1037
+ # with jax.disable_jit():
1038
+ # f(1.)
1039
+ #
1040
+ # @parameterized.product(
1041
+ # grad_fun=[brainstate.augment.grad, brainstate.augment.vector_grad]
1042
+ # )
1043
+ # def test_print_info1(self, grad_fun):
1044
+ # file = tempfile.TemporaryFile(mode='w+')
1045
+ #
1046
+ # @functools.partial(grad_fun, argnums=0)
1047
+ # def f2(a, b):
1048
+ # print('compiling f2 ...', file=file)
1049
+ # return a + b
1050
+ #
1051
+ # @functools.partial(grad_fun, argnums=0)
1052
+ # def f1(a):
1053
+ # print('compiling f1 ...', file=file)
1054
+ # return f2(a, 1.)
1055
+ #
1056
+ # expect_res = '''
1057
+ # compiling f1 ...
1058
+ # compiling f2 ...
1059
+ # compiling f1 ...
1060
+ # compiling f2 ...
1061
+ # '''
1062
+ #
1063
+ # print(f1(1.))
1064
+ # file.seek(0)
1065
+ # self.assertTrue(file.read().strip() == expect_res.strip())
1066
+ #
1067
+ # file = tempfile.TemporaryFile(mode='w+')
1068
+ # with jax.disable_jit():
1069
+ # expect_res = '''
1070
+ # compiling f1 ...
1071
+ # compiling f2 ...
1072
+ # '''
1073
+ # self.assertTrue(f1(1.) == 0.)
1074
+ # file.seek(0)
1075
+ # self.assertTrue(file.read().strip() == expect_res.strip())
1076
+ #
1077
+ # @parameterized.product(
1078
+ # grad_fun=[brainstate.augment.grad, brainstate.augment.vector_grad]
1079
+ # )
1080
+ # def test_print_info2(self, grad_fun):
1081
+ # file = tempfile.TemporaryFile(mode='w+')
1082
+ #
1083
+ # @functools.partial(grad_fun, argnums=0)
1084
+ # def f1(a):
1085
+ # @functools.partial(grad_fun, argnums=0)
1086
+ # def f2(a, b):
1087
+ # print('compiling f2 ...', file=file)
1088
+ # return a + b
1089
+ #
1090
+ # print('compiling f1 ...', file=file)
1091
+ # return f2(a, 1.)
1092
+ #
1093
+ # expect_res = '''
1094
+ # compiling f1 ...
1095
+ # compiling f2 ...
1096
+ # compiling f1 ...
1097
+ # compiling f2 ...
1098
+ # compiling f2 ...
1099
+ # '''
1100
+ # self.assertTrue(f1(1.) == 0.)
1101
+ # file.seek(0)
1102
+ # self.assertTrue(file.read().strip() == expect_res.strip())
1103
+ #
1104
+ # file = tempfile.TemporaryFile(mode='w+')
1105
+ # with jax.disable_jit():
1106
+ # expect_res = '''
1107
+ # compiling f1 ...
1108
+ # compiling f2 ...
1109
+ # '''
1110
+ # self.assertTrue(f1(1.) == 0.)
1111
+ # file.seek(0)
1112
+ # # print(file.read().strip())
1113
+ # self.assertTrue(file.read().strip() == expect_res.strip())
1114
+ #
1115
+ # def test_debug_correctness1(self):
1116
+ # def test_f():
1117
+ # a = jnp.Variable(jnp.ones(2))
1118
+ # b = jnp.Variable(jnp.zeros(2))
1119
+ #
1120
+ # @brainstate.augment.vector_grad(argnums=0)
1121
+ # def f1(c):
1122
+ # a.value += 1
1123
+ # b.value += 10
1124
+ # return a * b * c
1125
+ #
1126
+ # return a, b, f1(1.)
1127
+ #
1128
+ # r1 = test_f()
1129
+ # print(r1)
1130
+ #
1131
+ # with jax.disable_jit():
1132
+ # r2 = test_f()
1133
+ # print(r2)
1134
+ # self.assertTrue(jnp.allclose(r1[0], r2[0]))
1135
+ # self.assertTrue(jnp.allclose(r1[1], r2[1]))
1136
+ # self.assertTrue(jnp.allclose(r1[2], r2[2]))
1137
+ #
1138
+ # def f1(c, a, b):
1139
+ # a += 1
1140
+ # b += 10
1141
+ # return a * b * c
1142
+ #
1143
+ # r3 = vgrad(f1, 1., jnp.ones(2).value, jnp.zeros(2).value)
1144
+ # self.assertTrue(jnp.allclose(r1[2], r3))
1145
+ #
1146
+ # def _bench_f2(self, dd):
1147
+ # a = jnp.Variable(jnp.ones(2))
1148
+ # b = jnp.Variable(jnp.zeros(2))
1149
+ #
1150
+ #
1151
+ # def run_fun(d):
1152
+ # @brainstate.augment.vector_grad(argnums=0)
1153
+ # def f1(c):
1154
+ # a.value += d
1155
+ # b.value += 10
1156
+ # return a * b * c
1157
+ #
1158
+ # return a, b, f1(1.)
1159
+ #
1160
+ # return run_fun(dd)
1161
+ #
1162
+ # def test_debug_correctness2(self):
1163
+ # r1 = self._bench_f2(1.)
1164
+ # print(r1)
1165
+ #
1166
+ # with jax.disable_jit():
1167
+ # r2 = self._bench_f2(1.)
1168
+ # print(r2)
1169
+ #
1170
+ # self.assertTrue(jnp.allclose(r1[0], r2[0]))
1171
+ # self.assertTrue(jnp.allclose(r1[1], r2[1]))
1172
+ # self.assertTrue(jnp.allclose(r1[2], r2[2]))
1173
+ #
1174
+ # def test_cache1(self):
1175
+ # file = tempfile.TemporaryFile(mode='w+')
1176
+ #
1177
+ # def f(a, b):
1178
+ # print('compiling f ...', file=file)
1179
+ # return a + b
1180
+ #
1181
+ # grad1 = brainstate.augment.grad(f)(1., 2.) # call "f" twice, one for Variable finding, one for compiling
1182
+ # grad2 = brainstate.augment.vector_grad(f)(1., 2.) # call "f" once for compiling
1183
+ #
1184
+ # file.seek(0)
1185
+ # print(file.read().strip())
1186
+ #
1187
+ # expect_res = '''
1188
+ # compiling f ...
1189
+ # compiling f ...
1190
+ # compiling f ...
1191
+ # '''
1192
+ # file.seek(0)
1193
+ # self.assertTrue(file.read().strip() == expect_res.strip())
1194
+ #
1195
+ #
1196
+
1197
+
1198
+ class TestUnitAwareGrad(unittest.TestCase):
1199
+ def test_grad1(self):
1200
+ def f(x):
1201
+ return u.math.sum(x ** 2)
1202
+
1203
+ x = jnp.array([1., 2., 3.]) * u.ms
1204
+ g = brainstate.augment.grad(f, unit_aware=True)(x)
1205
+ self.assertTrue(u.math.allclose(g, 2 * x))
1206
+
1207
+ def test_vector_grad1(self):
1208
+ def f(x):
1209
+ return x ** 3
1210
+
1211
+ x = jnp.array([1., 2., 3.]) * u.ms
1212
+ g = brainstate.augment.vector_grad(f, unit_aware=True)(x)
1213
+ self.assertTrue(u.math.allclose(g, 3 * x ** 2))
1214
+
1215
+ def test_jacrev1(self):
1216
+ def f(x, y):
1217
+ return u.math.asarray([x[0] * y[0],
1218
+ 5 * x[2] * y[1],
1219
+ 4 * x[1] ** 2, ])
1220
+
1221
+ _x = jnp.array([1., 2., 3.]) * u.ms
1222
+ _y = jnp.array([10., 5.]) * u.ms
1223
+
1224
+ g = brainstate.augment.jacrev(f, unit_aware=True, argnums=(0, 1))(_x, _y)
1225
+ self.assertTrue(
1226
+ u.math.allclose(
1227
+ g[0],
1228
+ u.math.asarray([
1229
+ [10., 0., 0.],
1230
+ [0., 0., 25.],
1231
+ [0., 16., 0.]
1232
+ ]) * u.ms
1233
+ )
1234
+ )
1235
+
1236
+ self.assertTrue(
1237
+ u.math.allclose(
1238
+ g[1],
1239
+ u.math.asarray([
1240
+ [1., 0.],
1241
+ [0., 15.],
1242
+ [0., 0.]
1243
+ ]) * u.ms
1244
+ )
1245
+ )
1246
+
1247
+ def test_jacfwd1(self):
1248
+ def f(x, y):
1249
+ return u.math.asarray([x[0] * y[0],
1250
+ 5 * x[2] * y[1],
1251
+ 4 * x[1] ** 2, ])
1252
+
1253
+ _x = jnp.array([1., 2., 3.]) * u.ms
1254
+ _y = jnp.array([10., 5.]) * u.ms
1255
+
1256
+ g = brainstate.augment.jacfwd(f, unit_aware=True, argnums=(0, 1))(_x, _y)
1257
+ self.assertTrue(
1258
+ u.math.allclose(
1259
+ g[0],
1260
+ u.math.asarray([
1261
+ [10., 0., 0.],
1262
+ [0., 0., 25.],
1263
+ [0., 16., 0.]
1264
+ ]) * u.ms
1265
+ )
1266
+ )
1267
+
1268
+ self.assertTrue(
1269
+ u.math.allclose(
1270
+ g[1],
1271
+ u.math.asarray([
1272
+ [1., 0.],
1273
+ [0., 15.],
1274
+ [0., 0.]
1275
+ ]) * u.ms
1276
+ )
1277
+ )
1278
+
1279
+ def test_hessian(self):
1280
+ unit = u.ms
1281
+
1282
+ def scalar_function(x):
1283
+ return x ** 3 + 3 * x * unit * unit + 2 * unit * unit * unit
1284
+
1285
+ hess = brainstate.augment.hessian(scalar_function, unit_aware=True)
1286
+ x = jnp.array(1.0) * unit
1287
+ res = hess(x)
1288
+ expected_hessian = jnp.array([[6.0]]) * unit
1289
+ assert u.math.allclose(res, expected_hessian)