Source file owl_algodiff_ops_builder.ml
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
# 1 "src/base/algodiff/owl_algodiff_ops_builder.ml"
module Make (Core : Owl_algodiff_core_sig.Sig) = struct
open Core
let cmp_tag ai bi = if ai > bi then 1 else if ai < bi then -1 else 0
let op_siso ~ff ~fd ~df ~r a =
match a with
| DF (ap, at, ai) ->
let cp = fd ap in
DF (cp, df cp ap at, ai)
| DR (ap, _, _, _, ai, _) ->
let cp = fd ap in
DR (cp, ref (zero cp), r a, ref 0, ai, ref 0)
| ap -> ff ap
let op_sipo ~ff ~fd ~df ~r a =
match a with
| DF (ap, at, ai) ->
let cp1, cp2 = fd ap in
DF (cp1, df cp1 ap at, ai), DF (cp2, df cp2 ap at, ai)
| DR (ap, _, _, _, ai, _) ->
let cp1, cp2 = fd ap in
let ca1_ref = ref (zero cp1) in
let ca2_ref = ref (zero cp2) in
let cp1_ref = ref cp1 in
let cp2_ref = ref cp2 in
let tracker = ref 0 in
( DR
(cp1, ca1_ref, r (a, (cp1_ref, cp2_ref), (ca1_ref, ca2_ref)), ref 0, ai, tracker)
, DR
(cp2, ca2_ref, r (a, (cp1_ref, cp2_ref), (ca1_ref, ca2_ref)), ref 0, ai, tracker)
)
| ap -> ff ap
let op_sito ~ff ~fd ~df ~r a =
match a with
| DF (ap, at, ai) ->
let cp1, cp2, cp3 = fd ap in
DF (cp1, df cp1 ap at, ai), DF (cp2, df cp2 ap at, ai), DF (cp3, df cp3 ap at, ai)
| DR (ap, _, _, _, ai, _) ->
let cp1, cp2, cp3 = fd ap in
let ca1_ref = ref (zero cp1) in
let ca2_ref = ref (zero cp2) in
let ca3_ref = ref (zero cp3) in
let cp1_ref = ref cp1 in
let cp2_ref = ref cp2 in
let cp3_ref = ref cp3 in
let tracker = ref 0 in
( DR
( cp1
, ca1_ref
, r (a, (cp1_ref, cp2_ref, cp3_ref), (ca1_ref, ca2_ref, ca3_ref))
, ref 0
, ai
, tracker )
, DR
( cp2
, ca2_ref
, r (a, (cp1_ref, cp2_ref, cp3_ref), (ca1_ref, ca2_ref, ca3_ref))
, ref 0
, ai
, tracker )
, DR
( cp3
, ca3_ref
, r (a, (cp1_ref, cp2_ref, cp3_ref), (ca1_ref, ca2_ref, ca3_ref))
, ref 0
, ai
, tracker ) )
| ap -> ff ap
let op_siao ~ff ~fd ~df ~r a =
match a with
| DF (ap, at, ai) ->
let cp_arr = fd ap in
Array.map (fun cp -> DF (cp, df cp ap at, ai)) cp_arr
| DR (ap, _, _, _, ai, _) ->
let cp_arr = fd ap in
let cp_arr_ref = Array.map (fun cp -> ref cp) cp_arr in
let tracker = ref 0 in
let ca_ref_arr = Array.map (fun cp -> ref (zero cp)) cp_arr in
Array.map2
(fun cp ca_ref ->
DR (cp, ca_ref, r (a, cp_arr_ref, ca_ref_arr), ref 0, ai, tracker))
cp_arr
ca_ref_arr
| ap -> ff ap
let op_piso ~ff ~fd ~df_da ~df_db ~df_dab ~r_d_d ~r_d_c ~r_c_d a b =
match a, b with
| F _ap, DF (bp, bt, bi) ->
let cp = fd a bp in
DF (cp, df_db cp a bp bt, bi)
| DF (ap, at, ai), F _bp ->
let cp = fd ap b in
DF (cp, df_da cp ap at b, ai)
| Arr _ap, DF (bp, bt, bi) ->
let cp = fd a bp in
DF (cp, df_db cp a bp bt, bi)
| DF (ap, at, ai), Arr _bp ->
let cp = fd ap b in
DF (cp, df_da cp ap at b, ai)
| F _ap, DR (bp, _, _, _, bi, _) ->
let cp = fd a bp in
DR (cp, ref (zero cp), r_c_d a b, ref 0, bi, ref 0)
| DR (ap, _, _, _, ai, _), F _bp ->
let cp = fd ap b in
DR (cp, ref (zero cp), r_d_c a b, ref 0, ai, ref 0)
| Arr _ap, DR (bp, _, _, _, bi, _) ->
let cp = fd a bp in
DR (cp, ref (zero cp), r_c_d a b, ref 0, bi, ref 0)
| DR (ap, _, _, _, ai, _), Arr _bp ->
let cp = fd ap b in
DR (cp, ref (zero cp), r_d_c a b, ref 0, ai, ref 0)
| DF (ap, at, ai), DR (bp, _, _, _, bi, _) ->
(match cmp_tag ai bi with
| 1 ->
let cp = fd ap b in
DF (cp, df_da cp ap at b, ai)
| -1 ->
let cp = fd a bp in
DR (cp, ref (zero cp), r_c_d a b, ref 0, bi, ref 0)
| _ -> failwith "error: forward and reverse clash at the same level")
| DR (ap, _, _, _, ai, _), DF (bp, bt, bi) ->
(match cmp_tag ai bi with
| -1 ->
let cp = fd a bp in
DF (cp, df_db cp a bp bt, bi)
| 1 ->
let cp = fd ap b in
DR (cp, ref (zero cp), r_d_c a b, ref 0, ai, ref 0)
| _ -> failwith "error: forward and reverse clash at the same level")
| DF (ap, at, ai), DF (bp, bt, bi) ->
(match cmp_tag ai bi with
| 0 ->
let cp = fd ap bp in
DF (cp, df_dab cp ap at bp bt, ai)
| 1 ->
let cp = fd ap b in
DF (cp, df_da cp ap at b, ai)
| _ ->
let cp = fd a bp in
DF (cp, df_db cp a bp bt, bi))
| DR (ap, _, _, _, ai, _), DR (bp, _, _, _, bi, _) ->
(match cmp_tag ai bi with
| 0 ->
let cp = fd ap bp in
DR (cp, ref (zero cp), r_d_d a b, ref 0, ai, ref 0)
| 1 ->
let cp = fd ap b in
DR (cp, ref (zero cp), r_d_c a b, ref 0, ai, ref 0)
| _ ->
let cp = fd a bp in
DR (cp, ref (zero cp), r_c_d a b, ref 0, bi, ref 0))
| a, b -> ff a b
module type Siso = sig
val label : string
val ff_f : A.elt -> t
val ff_arr : A.arr -> t
val df : t -> t -> t -> t
val dr : t -> t -> t ref -> t
end
let build_siso (module S : Siso) =
let rec f a =
let open S in
let ff = function
| F a -> S.ff_f a
| Arr a -> S.ff_arr a
| _ -> error_uniop label a
in
let fd a = f a in
let r a =
let adjoint cp ca t = (S.dr a cp ca, a) :: t in
let register t = a :: t in
let label = S.label, [ a ] in
adjoint, register, label
in
op_siso ~ff ~fd ~df:S.df ~r a
in
f
module type Sipo = sig
val label : string
val ff_f : A.elt -> t * t
val ff_arr : A.arr -> t * t
val df : t -> t -> t -> t
val dr : t -> t -> t ref * t ref -> t ref * t ref -> t
end
let build_sipo (module S : Sipo) =
let rec f a =
let open S in
let ff = function
| F a -> S.ff_f a
| Arr a -> S.ff_arr a
| _ -> error_uniop label a
in
let fd = f in
let r (a, cp_ref, ca_ref) =
let adjoint cp _ca t = (S.dr a cp cp_ref ca_ref, a) :: t in
let register t = a :: t in
let label = S.label, [ a ] in
adjoint, register, label
in
op_sipo ~ff ~fd ~df ~r a
in
f
module type Sito = sig
val label : string
val ff_f : A.elt -> t * t * t
val ff_arr : A.arr -> t * t * t
val df : t -> t -> t -> t
val dr : t -> t -> t ref * t ref * t ref -> t ref * t ref * t ref -> t
end
let build_sito (module S : Sito) =
let rec f a =
let open S in
let ff = function
| F a -> S.ff_f a
| Arr a -> S.ff_arr a
| _ -> error_uniop label a
in
let fd = f in
let r (a, cp_ref, ca_ref) =
let adjoint cp _ca t = (S.dr a cp cp_ref ca_ref, a) :: t in
let register t = a :: t in
let label = S.label, [ a ] in
adjoint, register, label
in
op_sito ~ff ~fd ~df ~r a
in
f
module type Siao = sig
val label : string
val ff_f : A.elt -> t array
val ff_arr : A.arr -> t array
val df : t -> t -> t -> t
val dr : t -> t -> t ref array -> t ref array -> t
end
let build_siao (module S : Siao) =
let rec f a =
let open S in
let ff = function
| F a -> S.ff_f a
| Arr a -> S.ff_arr a
| _ -> error_uniop label a
in
let fd = f in
let r (a, cp_arr_ref, ca_arr_ref) =
let adjoint cp _ca_ref t = (S.dr a cp cp_arr_ref ca_arr_ref, a) :: t in
let register t = a :: t in
let label = S.label, [ a ] in
adjoint, register, label
in
op_siao ~ff ~fd ~df ~r a
in
f
module type Piso = sig
val label : string
val ff_aa : A.elt -> A.elt -> t
val ff_ab : A.elt -> A.arr -> t
val ff_ba : A.arr -> A.elt -> t
val ff_bb : A.arr -> A.arr -> t
val df_da : t -> t -> t -> t -> t
val df_db : t -> t -> t -> t -> t
val df_dab : t -> t -> t -> t -> t -> t
val dr_ab : t -> t -> t -> t ref -> t * t
val dr_a : t -> t -> t -> t ref -> t
val dr_b : t -> t -> t -> t ref -> t
end
let build_piso (module S : Piso) =
let rec f a b =
let ff a b =
match a, b with
| F a, F b -> S.ff_aa a b
| F a, Arr b -> S.ff_ab a b
| Arr a, F b -> S.ff_ba a b
| Arr a, Arr b -> S.ff_bb a b
| _ -> error_binop S.label a b
in
let fd = f in
let r_d_d a b =
let adjoint cp ca_ref t =
let abar, bbar = S.dr_ab a b cp ca_ref in
(abar, a) :: (bbar, b) :: t
in
let register t = a :: b :: t in
let label = S.label ^ "_d_d", [ a; b ] in
adjoint, register, label
in
let r_d_c a b =
let adjoint cp ca_ref t = (S.dr_a a b cp ca_ref, a) :: t in
let register t = a :: t in
let label = S.label ^ "_d_c", [ a; b ] in
adjoint, register, label
in
let r_c_d a b =
let adjoint cp ca_ref t = (S.dr_b a b cp ca_ref, b) :: t in
let register t = b :: t in
let label = S.label ^ "_c_d", [ a; b ] in
adjoint, register, label
in
op_piso
~ff
~fd
~df_da:S.df_da
~df_db:S.df_db
~df_dab:S.df_dab
~r_d_d
~r_d_c
~r_c_d
a
b
in
f
module type Aiso = sig
val label : string
val ff : t array -> t
val df : int list -> t -> t array -> t array -> t
val dr : int list -> t array -> t -> t ref -> (t * t) list
end
let build_aiso =
let build_info =
Array.fold_left
(fun (i, t, m, idxs) x ->
match m, x with
| _, F _ | _, Arr _ -> succ i, t, m, idxs
| `normal, DR (_, _, _, _, t', _) -> succ i, t', `reverse, [ i ]
| `forward, DR (_, _, _, _, t', _) ->
if t' > t
then succ i, t', `reverse, []
else if t' = t
then failwith "error: forward and reverse clash on the same level"
else succ i, t, `forward, idxs
| `reverse, DR (_, _, _, _, t', _) ->
if t' > t
then succ i, t', `reverse, []
else if t' = t
then succ i, t', `reverse, i :: idxs
else succ i, t, m, idxs
| `normal, DF (_, _, t') -> succ i, t', `forward, [ i ]
| `forward, DF (_, _, t') ->
if t' > t
then succ i, t', `forward, []
else if t' = t
then succ i, t', `forward, i :: idxs
else succ i, t, `forward, idxs
| `reverse, DF (_, _, t') ->
if t' > t
then succ i, t', `forward, []
else if t' = t
then failwith "error: forward and reverse clash on the same level"
else succ i, t, `reverse, idxs)
(0, -10000, `normal, [])
in
fun (module S : Aiso) ->
let rec f a =
let _, t, mode, idxs = build_info a in
let idxs = idxs |> List.rev in
match mode with
| `normal -> S.ff a
| `forward ->
let cp =
Array.map
(fun x ->
match x with
| DF (p, _, t') -> if t = t' then p else x
| x -> x)
a
|> f
in
let at =
let ap = a |> Array.map primal in
let at = a |> Array.map zero in
List.iter (fun k -> at.(k) <- tangent a.(k)) idxs;
S.df idxs cp ap at
in
DF (cp, at, t)
| `reverse ->
let cp =
Array.map
(fun x ->
match x with
| DR (p, _, _, _, t', _) -> if t = t' then p else x
| x -> x)
a
|> f
in
let adjoint cp ca t = List.append (S.dr idxs a cp ca) t in
let register t = List.fold_left (fun t i -> a.(i) :: t) t idxs in
let label = S.label, List.(map (fun i -> a.(i)) idxs) in
DR (cp, ref (zero cp), (adjoint, register, label), ref 0, t, ref 0)
in
f
end