Source file owl_algodiff_core.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
# 1 "src/base/algodiff/owl_algodiff_core.ml"
module Make (A : Owl_types_ndarray_algodiff.Sig) = struct
include Owl_algodiff_types.Make (A)
module A = A
let _global_tag = ref 0
let tag () =
_global_tag := !_global_tag + 1;
!_global_tag
let reset_zero = function
| F _ -> F A.(float_to_elt 0.)
| Arr ap ->
A.reset ap;
Arr ap
| _ -> failwith "error: reset_zero"
let primal = function
| DF (ap, _, _) -> ap
| DR (ap, _, _, _, _, _) -> ap
| ap -> ap
let rec primal' = function
| DF (ap, _, _) -> primal' ap
| DR (ap, _, _, _, _, _) -> primal' ap
| ap -> ap
let rec zero = function
| F _ -> F A.(float_to_elt 0.)
| Arr ap -> Arr A.(zeros (shape ap))
| DF (ap, _, _) -> ap |> primal' |> zero
| DR (ap, _, _, _, _, _) -> ap |> primal' |> zero
let tangent = function
| DF (_, at, _) -> at
| DR _ -> failwith "error: no tangent for DR"
| ap -> zero ap
let adjref = function
| DF _ -> failwith "error: no adjref for DF"
| DR (_, at, _, _, _, _) -> at
| ap -> ref (zero ap)
let adjval = function
| DF _ -> failwith "error: no adjval for DF"
| DR (_, at, _, _, _, _) -> !at
| ap -> zero ap
let shape x =
match primal' x with
| F _ -> [||]
| Arr ap -> A.shape ap
| _ -> failwith "error: AD.shape"
let rec is_float x =
match x with
| Arr _ -> false
| F _ -> true
| DF _ -> is_float (primal' x)
| DR _ -> is_float (primal' x)
let rec is_arr x =
match x with
| Arr _ -> false
| F _ -> true
| DF _ -> is_arr (primal' x)
| DR _ -> is_arr (primal' x)
let row_num x = (shape x).(0)
let col_num x = (shape x).(1)
let numel x =
match primal' x with
| Arr x -> A.numel x
| _ -> failwith "error: AD.numel"
let clip_by_value ~amin ~amax x =
match primal' x with
| Arr x -> Arr A.(clip_by_value ~amin ~amax x)
| _ -> failwith "error: AD.clip_by_value"
let clip_by_l2norm a x =
match primal' x with
| Arr x -> Arr A.(clip_by_l2norm a x)
| _ -> failwith "error: AD.clip_by_l2norm"
let copy_primal' x =
match primal' x with
| Arr ap -> Arr A.(copy ap)
| _ -> failwith "error: AD.copy"
let tile x reps =
match primal' x with
| Arr x -> Arr A.(tile x reps)
| _ -> failwith "error: AD.tile"
let repeat x reps =
match primal' x with
| Arr x -> Arr A.(repeat x reps)
| _ -> failwith "error: AD.repeat"
let pack_elt x = F x
let unpack_elt x =
match primal x with
| F x -> x
| _ -> failwith "error: AD.unpack_elt"
let pack_flt x = F A.(float_to_elt x)
let _f x = F A.(float_to_elt x)
let unpack_flt x =
match primal x with
| F x -> A.elt_to_float x
| _ -> failwith "error: AD.unpack_flt"
let pack_arr x = Arr x
let unpack_arr x =
match primal x with
| Arr x -> x
| _ -> failwith "error: AD.unpack_arr"
let deep_info x =
match primal' x with
| F a -> Printf.sprintf "F(%g)" A.(elt_to_float a)
| Arr a ->
Printf.sprintf "Arr(%s)" (A.shape a |> Owl_utils_array.to_string string_of_int)
| _ -> "you should not have reached here!"
let type_info x =
match x with
| F _a -> Printf.sprintf "[%s]" (deep_info x)
| DF (ap, _at, ai) -> Printf.sprintf "[DF tag:%i ap:%s]" ai (deep_info ap)
| DR (ap, _at, _ao, _af, ai, _) ->
Printf.sprintf "[DR tag:%i ap:%s]" ai (deep_info ap)
| _ -> Printf.sprintf "[%s]" (deep_info x)
let error_binop op a b =
let s0 = "#0:" ^ type_info a in
let s1 = "#1:" ^ type_info b in
failwith (op ^ " : " ^ s0 ^ ", " ^ s1)
let error_uniop op a =
let s = type_info a in
failwith (op ^ " : " ^ s)
end