diff --git a/README.md b/README.md index 7e7c124..efa5b94 100644 --- a/README.md +++ b/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) diff --git a/src/AutoGrad.jl b/src/AutoGrad.jl index f27c63d..14b4f1f 100644 --- a/src/AutoGrad.jl +++ b/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") diff --git a/src/highorder.jl b/src/highorder.jl new file mode 100644 index 0000000..d757a2b --- /dev/null +++ b/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 \ No newline at end of file diff --git a/test/highorder.jl b/test/highorder.jl index 71c52e9..b5892ae 100644 --- a/test/highorder.jl +++ b/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