package owl-base

  1. Overview
  2. Docs
package owl-base
Legend:
Page
Library
Module
Module type
Parameter
Class
Class type
Source

Source file owl_algodiff_generic.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
# 1 "src/base/algodiff/owl_algodiff_generic.ml"
(*
 * OWL - OCaml Scientific Computing
 * Copyright (c) 2016-2022 Liang Wang <liang@ocaml.xyz>

 * Core nested automatic differentiation algorithm and differentiation API
 * ported from DiffSharp (http://diffsharp.github.io), copyright (c) 2014-2016
 * National University of Ireland Maynooth (Atilim Gunes Baydin), 2014-2018
 * National University of Ireland Maynooth (Barak A. Pearlmutter
 * <barak@pearlmutter.net>), 2016-2018 University of Oxford (Atilim Gunes
 * Baydin <gunes@robots.ox.ac.uk>), 2017-2018 Microsoft Research Cambridge
 * (Don Syme <dsyme@microsoft.com>
 *)

(* Functor of making AD module of different precisions *)

module Make (A : Owl_types_ndarray_algodiff.Sig) = struct
  (* include functions in the Core module *)
  module Core = Owl_algodiff_core.Make (A)
  include Core

  (* include graph conversion functions *)
  include Owl_algodiff_graph_convert.Make (Core)

  (* instantiate operations *)
  module Ops = Owl_algodiff_ops.Make (Core)
  include Ops

  (* include core reverse mode functions *)
  module Reverse = Owl_algodiff_reverse.Make (struct
    include Core

    let reverse_add = Maths.add
  end)

  (* convenient wrappers *)

  let make_forward p t i = DF (p, t, i)

  let make_reverse p i =
    let adjoint _cp _ca t = t in
    let register t = t in
    let label = "Noop", [] in
    DR (p, ref (zero p), (adjoint, register, label), ref 0, i, ref 0)


  (* expose reverse prop: propagate gradients *)
  let reverse_prop = Reverse.reverse_prop

  (* derivative of f (scalar -> scalar) at x, forward ad *)
  let diff' f x =
    if not (is_float x) then failwith "input of `diff` must be a scalar";
    let x = make_forward x (pack_flt 1.) (tag ()) in
    let y = f x in
    primal y, tangent y


  (* derivative of f (scalar -> scalar) at x, forward ad *)
  let diff f x = diff' f x |> snd

  (* gradient of f (vector -> scalar) at x, reverse ad *)
  let grad' f x =
    let x = make_reverse x (tag ()) in
    let y = f x in
    if not (is_float y) then failwith "output of `grad` must be a scalar";
    Reverse.reverse_reset y;
    Reverse.reverse_push (pack_flt 1.) y;
    primal y, x |> adjval


  (* gradient of f (vector -> scalar) at x, reverse ad *)
  let grad f x = grad' f x |> snd

  (* jacobian vector product of f (vector -> vector) at x along v, forward ad *)
  let jacobianv' f x v =
    if shape x <> shape v
    then failwith "jacobianv': vector not the same dimension as input";
    let x = make_forward x v (tag ()) in
    let y = f x in
    primal y, tangent y


  (* jacobian vector product of f (vector -> vector) at x along v, forward ad *)
  let jacobianv f x v = jacobianv' f x v |> snd

  (* transposed jacobian vector product of f (vector -> vector) at x along v, backward ad *)
  let jacobianTv' f x v =
    let x = make_reverse x (tag ()) in
    let y = f x in
    if shape y <> shape v
    then failwith "jacobianTv': vector not the same dimension as output";
    Reverse.reverse_reset y;
    Reverse.reverse_push v y;
    primal y, x |> adjval


  (* transposed jacobian vector product of f (vector -> vector) at x along v, backward ad *)
  let jacobianTv f x v = jacobianTv' f x v |> snd

  (* jacobian of f (vector -> vector) at x, both x and y are row vectors, also return the
     original value. The jacobian J satisfies dx = J df *)
  let jacobian' =
    let dim_typ x =
      match primal' x with
      | F _   -> `float
      | Arr x ->
        let s = A.shape x in
        let d = Array.length s in
        if d > 2
        then `arr
        else if s.(0) = 1
        then `row s.(1)
        else if s.(1) = 1
        then `col s.(0)
        else `mat
      | _     -> assert false
    in
    fun f x ->
      let y = f x in
      let m, n =
        match dim_typ y, dim_typ x with
        | `row a, `row 1 -> a, 1
        | `row a, `row b -> a, b
        | _              -> failwith "jacobian: input and output must both be row vectors"
      in
      let z =
        let jvps =
          if m > n
          then
            Array.init n (fun i ->
                let v = A.zeros [| 1; n |] in
                A.(set v [| 0; i |] (float_to_elt 1.));
                jacobianv f x (Arr v))
          else
            Array.init m (fun i ->
                let v = A.zeros [| 1; m |] in
                A.(set v [| 0; i |] (float_to_elt 1.));
                jacobianTv f x (Arr v))
        in
        match y with
        | F _         -> failwith "input to jacobian must be a row vector not a float"
        | Arr _       ->
          let z = A.zeros [| n; m |] in
          jvps
          |> Array.iteri (fun i v ->
                 match v with
                 | Arr v ->
                   if m > n
                   then A.copy_row_to v z i
                   else A.copy_col_to (A.transpose v) z i
                 | _     -> failwith "error: jacobian");
          Arr z
        | DF _ | DR _ ->
          if m > n
          then Ops.Maths.concatenate ~axis:0 jvps
          else Ops.Maths.concatenate ~axis:0 jvps |> Ops.Maths.transpose
      in
      primal y, z


  (* jacobian of f *)
  let jacobian f x = jacobian' f x |> snd

  (* gradient, hessian of f (vector -> scalar) at [x] *)
  let gradhessian f x = jacobian' (grad f) x

  (* original value, gradient, and hessian *)
  let gradhessian' f x =
    let g, h = gradhessian f x in
    f x, g, h


  (* hessian of f *)
  let hessian f x = (f |> grad |> jacobian) x

  (* original value and hessian of f *)
  let hessian' f x = f x, hessian f x

  (* original value, gradient-vector product, hessian-vector product *)
  let gradhessianv' f x v =
    let gv, hv = grad' (fun y -> jacobianv f y v) x in
    f x, gv, hv


  (* gradient-vector product and hessian-vector product *)
  let gradhessianv f x v =
    let _, gv, hv = gradhessianv' f x v in
    gv, hv


  (* original value and hessian-vector product *)
  let hessianv' f x v =
    let fv, _, hv = gradhessianv' f x v in
    fv, hv


  (* hessian-vector *)
  let hessianv f x v =
    let _, _, hv = gradhessianv' f x v in
    hv


  (* laplacian of f *)
  let laplacian f x = F (hessian f x |> unpack_arr |> A.trace)

  let laplacian' f x = f x, laplacian f x
end

(* ends here *)