brainstate 0.2.1__py2.py3-none-any.whl → 0.2.2__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 (115) hide show
  1. brainstate/__init__.py +167 -169
  2. brainstate/_compatible_import.py +340 -340
  3. brainstate/_compatible_import_test.py +681 -681
  4. brainstate/_deprecation.py +210 -210
  5. brainstate/_deprecation_test.py +2297 -2319
  6. brainstate/_error.py +45 -45
  7. brainstate/_state.py +2157 -1652
  8. brainstate/_state_test.py +1129 -52
  9. brainstate/_utils.py +47 -47
  10. brainstate/environ.py +1495 -1495
  11. brainstate/environ_test.py +1223 -1223
  12. brainstate/graph/__init__.py +22 -22
  13. brainstate/graph/_node.py +240 -240
  14. brainstate/graph/_node_test.py +589 -589
  15. brainstate/graph/_operation.py +1620 -1624
  16. brainstate/graph/_operation_test.py +1147 -1147
  17. brainstate/mixin.py +1447 -1433
  18. brainstate/mixin_test.py +1017 -1017
  19. brainstate/nn/__init__.py +146 -137
  20. brainstate/nn/_activations.py +1100 -1100
  21. brainstate/nn/_activations_test.py +354 -354
  22. brainstate/nn/_collective_ops.py +635 -633
  23. brainstate/nn/_collective_ops_test.py +774 -774
  24. brainstate/nn/_common.py +226 -226
  25. brainstate/nn/_common_test.py +134 -154
  26. brainstate/nn/_conv.py +2010 -2010
  27. brainstate/nn/_conv_test.py +849 -849
  28. brainstate/nn/_delay.py +575 -575
  29. brainstate/nn/_delay_test.py +243 -243
  30. brainstate/nn/_dropout.py +618 -618
  31. brainstate/nn/_dropout_test.py +480 -477
  32. brainstate/nn/_dynamics.py +870 -1267
  33. brainstate/nn/_dynamics_test.py +53 -67
  34. brainstate/nn/_elementwise.py +1298 -1298
  35. brainstate/nn/_elementwise_test.py +829 -829
  36. brainstate/nn/_embedding.py +408 -408
  37. brainstate/nn/_embedding_test.py +156 -156
  38. brainstate/nn/_event_fixedprob.py +233 -233
  39. brainstate/nn/_event_fixedprob_test.py +115 -115
  40. brainstate/nn/_event_linear.py +83 -83
  41. brainstate/nn/_event_linear_test.py +121 -121
  42. brainstate/nn/_exp_euler.py +254 -254
  43. brainstate/nn/_exp_euler_test.py +377 -377
  44. brainstate/nn/_linear.py +744 -744
  45. brainstate/nn/_linear_test.py +475 -475
  46. brainstate/nn/_metrics.py +1070 -1070
  47. brainstate/nn/_metrics_test.py +611 -611
  48. brainstate/nn/_module.py +391 -384
  49. brainstate/nn/_module_test.py +40 -40
  50. brainstate/nn/_normalizations.py +1334 -1334
  51. brainstate/nn/_normalizations_test.py +699 -699
  52. brainstate/nn/_paddings.py +1020 -1020
  53. brainstate/nn/_paddings_test.py +722 -722
  54. brainstate/nn/_poolings.py +2239 -2239
  55. brainstate/nn/_poolings_test.py +952 -952
  56. brainstate/nn/_rnns.py +946 -946
  57. brainstate/nn/_rnns_test.py +592 -592
  58. brainstate/nn/_utils.py +216 -216
  59. brainstate/nn/_utils_test.py +401 -401
  60. brainstate/nn/init.py +809 -809
  61. brainstate/nn/init_test.py +180 -180
  62. brainstate/random/__init__.py +270 -270
  63. brainstate/random/{_rand_funs.py → _fun.py} +3938 -3938
  64. brainstate/random/{_rand_funs_test.py → _fun_test.py} +638 -640
  65. brainstate/random/_impl.py +672 -0
  66. brainstate/random/{_rand_seed.py → _seed.py} +675 -675
  67. brainstate/random/{_rand_seed_test.py → _seed_test.py} +48 -48
  68. brainstate/random/{_rand_state.py → _state.py} +1320 -1617
  69. brainstate/random/{_rand_state_test.py → _state_test.py} +551 -551
  70. brainstate/transform/__init__.py +56 -59
  71. brainstate/transform/_ad_checkpoint.py +176 -176
  72. brainstate/transform/_ad_checkpoint_test.py +49 -49
  73. brainstate/transform/_autograd.py +1025 -1025
  74. brainstate/transform/_autograd_test.py +1289 -1289
  75. brainstate/transform/_conditions.py +316 -316
  76. brainstate/transform/_conditions_test.py +220 -220
  77. brainstate/transform/_error_if.py +94 -94
  78. brainstate/transform/_error_if_test.py +52 -52
  79. brainstate/transform/_find_state.py +200 -0
  80. brainstate/transform/_find_state_test.py +84 -0
  81. brainstate/transform/_jit.py +399 -399
  82. brainstate/transform/_jit_test.py +143 -143
  83. brainstate/transform/_loop_collect_return.py +675 -675
  84. brainstate/transform/_loop_collect_return_test.py +58 -58
  85. brainstate/transform/_loop_no_collection.py +283 -283
  86. brainstate/transform/_loop_no_collection_test.py +50 -50
  87. brainstate/transform/_make_jaxpr.py +2176 -2016
  88. brainstate/transform/_make_jaxpr_test.py +1634 -1510
  89. brainstate/transform/_mapping.py +607 -529
  90. brainstate/transform/_mapping_test.py +104 -194
  91. brainstate/transform/_progress_bar.py +255 -255
  92. brainstate/transform/_unvmap.py +256 -256
  93. brainstate/transform/_util.py +286 -286
  94. brainstate/typing.py +837 -837
  95. brainstate/typing_test.py +780 -780
  96. brainstate/util/__init__.py +27 -27
  97. brainstate/util/_others.py +1024 -1024
  98. brainstate/util/_others_test.py +962 -962
  99. brainstate/util/_pretty_pytree.py +1301 -1301
  100. brainstate/util/_pretty_pytree_test.py +675 -675
  101. brainstate/util/_pretty_repr.py +462 -462
  102. brainstate/util/_pretty_repr_test.py +696 -696
  103. brainstate/util/filter.py +945 -945
  104. brainstate/util/filter_test.py +911 -911
  105. brainstate/util/struct.py +910 -910
  106. brainstate/util/struct_test.py +602 -602
  107. {brainstate-0.2.1.dist-info → brainstate-0.2.2.dist-info}/METADATA +108 -108
  108. brainstate-0.2.2.dist-info/RECORD +111 -0
  109. {brainstate-0.2.1.dist-info → brainstate-0.2.2.dist-info}/licenses/LICENSE +202 -202
  110. brainstate/transform/_eval_shape.py +0 -145
  111. brainstate/transform/_eval_shape_test.py +0 -38
  112. brainstate/transform/_random.py +0 -171
  113. brainstate-0.2.1.dist-info/RECORD +0 -111
  114. {brainstate-0.2.1.dist-info → brainstate-0.2.2.dist-info}/WHEEL +0 -0
  115. {brainstate-0.2.1.dist-info → brainstate-0.2.2.dist-info}/top_level.txt +0 -0
@@ -1,1289 +1,1289 @@
1
- # Copyright 2024 BrainX Ecosystem Limited. All Rights Reserved.
2
- #
3
- # Licensed under the Apache License, Version 2.0 (the "License");
4
- # you may not use this file except in compliance with the License.
5
- # You may obtain a copy of the License at
6
- #
7
- # http://www.apache.org/licenses/LICENSE-2.0
8
- #
9
- # Unless required by applicable law or agreed to in writing, software
10
- # distributed under the License is distributed on an "AS IS" BASIS,
11
- # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
- # See the License for the specific language governing permissions and
13
- # limitations under the License.
14
- # ==============================================================================
15
-
16
- # -*- 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.transform._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 BrainX Ecosystem Limited. All Rights Reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+ # ==============================================================================
15
+
16
+ # -*- 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.transform._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)