File tree Expand file tree Collapse file tree 1 file changed +28
-20
lines changed
Expand file tree Collapse file tree 1 file changed +28
-20
lines changed Original file line number Diff line number Diff line change 2020output = [0.0 ]
2121doutput = [1.0 ]
2222
23- fetch (schedule (Task (64 * 1024 ^ 2 ) do
24- autodiff (
25- Reverse,
26- f,
27- Const (tree),
28- Duplicated (X, dX),
29- Const (operators),
30- Duplicated (output, doutput),
31- )
32- end ))
23+ fetch (
24+ schedule (
25+ Task (64 * 1024 ^ 2 ) do
26+ autodiff (
27+ Reverse,
28+ f,
29+ Const (tree),
30+ Duplicated (X, dX),
31+ Const (operators),
32+ Duplicated (output, doutput),
33+ )
34+ end ,
35+ ),
36+ )
3337
3438true_dX = cat (ones (100 ), - sin .(X[2 , :]), zeros (100 ); dims= 2 )'
3539
@@ -50,16 +54,20 @@ d_tree = begin
5054 storage_tree = copy (tree)
5155 # Set all constants to zero:
5256 Enzyme. make_zero! (storage_tree)
53- fetch (schedule (Task (64 * 1024 ^ 2 ) do
54- autodiff (
55- Reverse,
56- my_loss_function,
57- Active,
58- Duplicated (tree, storage_tree),
59- Const (X),
60- Const (operators),
61- )
62- end ))
57+ fetch (
58+ schedule (
59+ Task (64 * 1024 ^ 2 ) do
60+ autodiff (
61+ Reverse,
62+ my_loss_function,
63+ Active,
64+ Duplicated (tree, storage_tree),
65+ Const (X),
66+ Const (operators),
67+ )
68+ end ,
69+ ),
70+ )
6371 storage_tree
6472end
6573
You can’t perform that action at this time.
0 commit comments