[wip] hessian, jacobian and their vector products

README.md

@@ -28,12 +28,12 @@ for a description of how the code works in detail.
 ## Installation
 
 You can install AutoGrad in Julia using:
-```
+```julia
 julia> Pkg.add("AutoGrad")
 ```
 
 In order to use it in your code start with:
-```
+```julia
 using AutoGrad
 ```
 
@@ -42,7 +42,7 @@ using AutoGrad
 Here is a linear regression example simplified from
 [housing.jl](https://github.com/denizyuret/AutoGrad.jl/blob/master/examples/housing.jl):
 
-```
+```julia
 using AutoGrad
 
 function loss(w)
@@ -94,7 +94,7 @@ or using the `@primitive` and `@zerograd` macros in
 [util.jl](https://github.com/denizyuret/AutoGrad.jl/blob/master/src/util.jl)
 Here is an example:
 
-```
+```julia
 @primitive hypot(x1::Number,x2::Number),dy,y  (dy*x1/y)  (dy*x2/y)
 ```
 
@@ -111,6 +111,52 @@ example `hypot(x1::Array,x2::Array)` is another hypot method.  In
 AutoGrad.jl each method can independently be defined as a primitive
 and can have its own specific gradient.
 
+## Jacobian and Higher Order Derivatives
+Since `grad` typically returns a differentiable function, it is really easy
+to take higher order derivatives of functions with scalar input and scalar output:
+```julia
+julia> g1 = grad(sin)
+(::gradfun) (generic function with 1 method)
+
+julia> g1(1) == cos(1)
+true
+
+julia> g2 = grad(g1)
+(::gradfun) (generic function with 1 method)
+
+julia> g2(1) == -sin(1)
+true   
+```
+More care has to be taken when functions with vector inputs are considered instead, 
+since the resulting gradient will be vector valued and AutoGrad is only able to
+differentiate scalar valued functions. In this case each output component of the gradient
+has to be differentiated separately in order to obtain the Hessian.
+AutoGrad provides the utilty function `hessian`, such that
+
+```julia
+julia> A = rand(3, 3);
+
+julia> f(x) = x'*A*x/2;
+
+julia> hessian(f)(rand(3)) == (A + A')/2
+true
+```
+Similarly the `jacobian` method returns the Jacobian matrix
+of a vector valued function:
+```julia
+
+```julia
+julia> A = rand(2, 3);
+
+julia> f(x) = A*x;
+
+julia> jacobian(f)(rand(3)) == (A + A')/2
+true
+```
+Computing Jacobians and Hessians can be heavily resource expensive, 
+both in time and in memory. In some use cases Jacobian/Hessian product 
+operators `jvp,vjp,hvp,vhp` can be used as a much faster alternative. 
+
 ## Code structure
 
 [core.jl](https://github.com/denizyuret/AutoGrad.jl/blob/master/src/core.jl)

src/AutoGrad.jl

@@ -21,6 +21,7 @@ end
 importall Base  # defining getindex, sin, etc.
 export grad, gradloss, check_grads, gradcheck, gradcheckN, getval
 export @primitive, @zerograd, recorder, Rec, Grad  # the last three are required for the macros to work
+export jacobian, hessian, jvp, vjp, hvp, vhp
 datapath = joinpath(dirname(@__FILE__),"..","data")
 
 include("core.jl")
@@ -29,6 +30,7 @@ include("dotfunc.jl")
 include("gradcheck.jl")
 include("util.jl")
 include("interfaces.jl")
+include("highorder.jl")
 include("base/reduce.jl")
 include("base/number.jl")
 include("base/float.jl")

src/highorder.jl

@@ -0,0 +1,77 @@
+"""
+    jacobian(f)
+
+Computes the Jacobian of the vector valued function `f`. 
+For `n = length(x)`, `m = length(f(x))`, the Jacobian
+matrix  `jacobian(f)(x)` has dimension `m x n`.
+"""
+function jacobian(f)
+    x -> begin
+        m = length(f(x))
+        hcat([grad(x->f(x)[i])(x) for i=1:m]...)'
+    end
+end
+
+"""
+    vjp(f)
+
+Returns the  vector-Jacobian product operator
+for the vector valued function `f`: 
+
+    vjp(f)(x, v) == v' * jacobian(f)(x) 
+
+This is more efficient than computing the full Jacobian
+if only a few products are needed. See also `jvp`.
+"""
+function vjp(f)
+    (x, v) -> begin
+        grad(x -> v'f(x))(x)'
+    end
+end
+
+
+
+# ref https://j-towns.github.io/2017/06/12/A-new-trick.html
+"""
+    jvp(f)
+
+Returns the  Jacobian-vector product operator
+for the vector valued function `f`: 
+
+    jvp(f)(x, u) == jacobian(f)(x) * u
+
+This is more efficient than computing the full Jacobian
+if only a few products are needed. See also `vjp`.
+"""
+function jvp(f)
+    (x, u) -> begin
+        v = f(x) # any value of v should do
+        g = v -> vjp(f)(x, v)'
+        vjp(g)(v, u)'        
+    end
+end
+
+"""
+    hessian(f)
+
+Compute the hessian of the scalar valued function `f`. 
+It is equivalent to `jacobian(grad(f))`.
+"""
+hessian(f) = jacobian(grad(f))
+
+"""
+    vhp(f)
+
+Returns the vector-Hessian product operator.
+We have the equivalence `vhp(f)(x, v) == vjp(grad(f))(x, v)`.
+"""
+vhp(f) = vjp(x->vec(grad(f)(x)))
+
+"""
+    hvp(f)
+
+Returns the Hessian product operat.
+It is equivalent to `jvp(jacobian(f))` 
+"""
+hvp(f) = (x,v) -> vhp(f)(x,v)'   # can use vjp instead of jvp since the 
+                        # hessian is symmetric

test/highorder.jl

@@ -10,6 +10,33 @@ include("header.jl")
     g7 = grad(g6);  @test g7(1)==-cos(1)
     g8 = grad(g7);  @test g8(1)==sin(1)
     g9 = grad(g8);  @test g9(1)==cos(1)
+
+    sym(A) = (A + A')/2
+
+    A = sym(rand(3, 3))
+    f(x) = x'*A*x/2
+    @test hessian(f)(rand(3)) == A
+    v = rand(3)
+    @test hvp(f)(rand(3), v) ≈ A*v
+    u = rand(3)
+    @test vhp(f)(rand(3), u) ≈ u'A
+
+    A = rand(2, 3)
+    f(x) = A*x
+    @test jacobian(f)(rand(3)) ≈ A 
+
+    A = rand(3, 3)
+    f(x) = x'*A*x/2
+    a = rand(3)
+    @test jacobian(f)(a) ≈ a' * sym(A)
+    @test hessian(f)(a) == sym(A)
+
+    A = rand(2, 3)
+    f(x) = A*x
+    v = rand(3)
+    @test jvp(f)(rand(3), v) == A*v
+    u = rand(2)
+    @test vjp(f)(rand(3), u) == u'A
 end
 
 nothing