Right preparation for second order closures (#122)

* Right preparation for second order closures

* Fix

Author: gdalle

ext/DifferentiationInterfaceFastDifferentiationExt/DifferentiationInterfaceFastDifferentiationExt.jl

@@ -6,11 +6,14 @@ import DifferentiationInterface as DI
 using FastDifferentiation:
     derivative,
     hessian,
+    hessian_times_v,
     jacobian,
     jacobian_times_v,
     jacobian_transpose_v,
     make_function,
-    make_variables
+    make_variables,
+    sparse_hessian,
+    sparse_jacobian
 using LinearAlgebra: dot
 using FastDifferentiation.RuntimeGeneratedFunctions: RuntimeGeneratedFunction
 
@@ -24,6 +27,9 @@ DI.supports_mutation(::AnyAutoFastDifferentiation) = DI.MutationNotSupported()
 myvec(x::Number) = [x]
 myvec(x::AbstractArray) = vec(x)
 
+issparse(::AutoFastDifferentiation) = false
+issparse(::AutoSparseFastDifferentiation) = true
+
 include("allocating.jl")
 
 end

ext/DifferentiationInterfaceFastDifferentiationExt/allocating.jl

@@ -64,12 +64,43 @@ end
 
 =#
 
+## Derivative
+
+function DI.prepare_derivative(f, ::AnyAutoFastDifferentiation, x)
+    x_var = only(make_variables(:x))
+    y_var = f(x_var)
+
+    x_vec_var = [x_var]
+    y_vec_var = y_var isa Number ? [y_var] : vec(y_var)
+    der_vec_var = derivative(y_vec_var, x_var)
+    der_exe = make_function(der_vec_var, x_vec_var; in_place=false)
+    return der_exe
+end
+
+function DI.value_and_derivative(f, ::AnyAutoFastDifferentiation, x, der_exe)
+    y = f(x)
+    der_vec = der_exe([x])
+    if y isa Number
+        return y, only(der_vec)
+    else
+        return y, reshape(der_vec, size(y))
+    end
+end
+
+function DI.value_and_derivative!!(f, der, backend::AnyAutoFastDifferentiation, x, der_exe)
+    return DI.value_and_derivative(f, backend, x, der_exe)
+end
+
 ## Jacobian
 
-function DI.prepare_jacobian(f, ::AnyAutoFastDifferentiation, x)
+function DI.prepare_jacobian(f, backend::AnyAutoFastDifferentiation, x)
     x_vec_var = make_variables(:x, size(x)...)
     y_vec_var = f(x_vec_var)
-    jac_var = jacobian(vec(y_vec_var), vec(x_vec_var))
+    if issparse(backend)
+        jac_var = sparse_jacobian(vec(y_vec_var), vec(x_vec_var))
+    else
+        jac_var = jacobian(vec(y_vec_var), vec(x_vec_var))
+    end
     jac_exe = make_function(jac_var, vec(x_vec_var); in_place=false)
     return jac_exe
 end
@@ -80,24 +111,85 @@ function DI.jacobian(
     return jac_exe(vec(x))
 end
 
-function DI.value_and_jacobian(f, backend, x, extras)
-    return f(x), DI.jacobian(f, backend, x, extras)
+function DI.value_and_jacobian(f, backend::AnyAutoFastDifferentiation, x, jac_exe)
+    return f(x), DI.jacobian(f, backend, x, jac_exe)
+end
+
+function DI.jacobian!!(f, jac, backend::AnyAutoFastDifferentiation, x, jac_exe)
+    return DI.jacobian(f, backend, x, jac_exe)
+end
+
+function DI.value_and_jacobian!!(f, jac, backend::AnyAutoFastDifferentiation, x, jac_exe)
+    return DI.value_and_jacobian(f, backend, x, jac_exe)
+end
+
+## Second derivative
+
+function DI.prepare_second_derivative(f, ::AnyAutoFastDifferentiation, x)
+    x_var = only(make_variables(:x))
+    y_var = f(x_var)
+
+    x_vec_var = [x_var]
+    y_vec_var = y_var isa Number ? [y_var] : vec(y_var)
+    der2_vec_var = derivative(y_vec_var, x_var, x_var)
+    der2_exe = make_function(der2_vec_var, x_vec_var; in_place=false)
+    return der2_exe
+end
+
+function DI.second_derivative(f, ::AnyAutoFastDifferentiation, x, der2_exe)
+    y = f(x)
+    der2_vec = der2_exe([x])
+    if y isa Number
+        return only(der2_vec)
+    else
+        return reshape(der2_vec, size(y))
+    end
+end
+
+function DI.second_derivative!!(f, der2, backend::AnyAutoFastDifferentiation, x, der2_exe)
+    return DI.second_derivative(f, backend, x, der2_exe)
+end
+
+## HVP
+
+function DI.prepare_hvp(f, ::AnyAutoFastDifferentiation, x)
+    x_var = if x isa Number
+        only(make_variables(:x))
+    else
+        make_variables(:x, size(x)...)
+    end
+    y_var = f(x_var)
+
+    x_vec_var = x_var isa Number ? [x_var] : vec(x_var)
+    hv_vec_var, v_vec_var = hessian_times_v(y_var, x_vec_var)
+    hvp_exe = make_function(hv_vec_var, [x_vec_var; v_vec_var]; in_place=false)
+    return hvp_exe
 end
 
-function DI.jacobian!!(f, backend::AnyAutoFastDifferentiation, x, extras)
-    return DI.jacobian(f, backend, x, extras)
+function DI.hvp(f, ::AnyAutoFastDifferentiation, x, v, hvp_exe::RuntimeGeneratedFunction)
+    v_vec = vcat(myvec(x), myvec(v))
+    hv_vec = hvp_exe(v_vec)
+    if x isa Number
+        return only(hv_vec)
+    else
+        return reshape(hv_vec, size(x))
+    end
 end
 
-function DI.value_and_jacobian!!(f, backend::AnyAutoFastDifferentiation, x, extras)
-    return DI.value_and_jacobian(f, backend, x, extras)
+function DI.hvp!!(f, p, backend::AnyAutoFastDifferentiation, x, v, hvp_exe)
+    return DI.hvp(f, backend, x, v, hvp_exe)
 end
 
 ## Hessian
 
-function DI.prepare_hessian(f, ::AnyAutoFastDifferentiation, x)
+function DI.prepare_hessian(f, backend::AnyAutoFastDifferentiation, x)
     x_vec_var = make_variables(:x, size(x)...)
     y_vec_var = f(x_vec_var)
-    hess_var = hessian(y_vec_var, vec(x_vec_var))
+    if issparse(backend)
+        hess_var = sparse_hessian(y_vec_var, vec(x_vec_var))
+    else
+        hess_var = hessian(y_vec_var, vec(x_vec_var))
+    end
     hess_exe = make_function(hess_var, vec(x_vec_var); in_place=false)
     return hess_exe
 end
@@ -108,11 +200,6 @@ function DI.hessian(
     return hess_exe(vec(x))
 end
 
-function DI.hessian(f, backend::AnyAutoFastDifferentiation, x, extras::Nothing)
-    hess_exe = prepare_hessian(f, backend, x)
+function DI.hessian!!(f, hess, backend::AnyAutoFastDifferentiation, x, hess_exe)
     return DI.hessian(f, backend, x, hess_exe)
 end
-
-function DI.hessian!!(f, backend::AnyAutoFastDifferentiation, x, extras)
-    return DI.hessian(f, backend, x, extras)
-end

ext/DifferentiationInterfaceZygoteExt/DifferentiationInterfaceZygoteExt.jl

@@ -61,4 +61,8 @@ function DI.hessian(f, ::AnyAutoZygote, x, extras::Nothing)
     return hessian(f, x)
 end
 
+function DI.hessian!!(f, hess, backend::AnyAutoZygote, x, extras::Nothing)
+    return DI.hessian(f, backend, x, extras)
+end
+
 end

lib/DifferentiationInterfaceTest/src/utils/printing.jl

@@ -16,6 +16,7 @@ pretty(::AutoTapir) = "Tapir"
 pretty(::AutoTracker) = "Tracker"
 pretty(::AutoZygote) = "Zygote"
 
+pretty(::AutoSparseFastDifferentiation) = "FastDifferentiation sparse"
 pretty(::AutoSparseFiniteDiff) = "FiniteDiff sparse"
 pretty(::AutoSparseForwardDiff) = "ForwardDiff sparse"
 pretty(::AutoSparsePolyesterForwardDiff) = "PolyesterForwardDiff sparse"

src/hessian.jl

@@ -4,7 +4,7 @@
     hessian(f, backend, x, [extras]) -> hess
 """
 function hessian(f, backend::AbstractADType, x, extras=prepare_hessian(f, backend, x))
-    new_backend = SecondOrder(backend, backend)
+    new_backend = SecondOrder(backend)
     new_extras = prepare_hessian(f, new_backend, x)
     return hessian(f, new_backend, x, new_extras)
 end
@@ -24,7 +24,7 @@ end
 function hessian!!(
     f, hess, backend::AbstractADType, x, extras=prepare_hessian(f, backend, x)
 )
-    new_backend = SecondOrder(backend, backend)
+    new_backend = SecondOrder(backend)
     new_extras = prepare_hessian(f, new_backend, x)
     return hessian!!(f, hess, new_backend, x, new_extras)
 end

src/hvp.jl

@@ -14,7 +14,7 @@ By order of preference:
     hvp(f, backend, x, v, [extras]) -> p
 """
 function hvp(f, backend::AbstractADType, x, v, extras=prepare_hvp(f, backend, x))
-    new_backend = SecondOrder(backend, backend)
+    new_backend = SecondOrder(backend)
     new_extras = prepare_hvp(f, new_backend, x)
     return hvp(f, new_backend, x, v, new_extras)
 end
@@ -32,26 +32,32 @@ end
 
 function hvp_aux(f, backend, x, v, extras, ::ForwardOverReverse)
     # JVP of the gradient
-    inner_extras = prepare_gradient(extras, f, inner(backend), x)
-    gradient_closure(z) = gradient(f, inner(backend), z, inner_extras)
+    function gradient_closure(z)
+        inner_extras = prepare_gradient(extras, f, inner(backend), z)
+        return gradient(f, inner(backend), z, inner_extras)
+    end
     outer_extras = prepare_pushforward(extras, gradient_closure, outer(backend), x)
     p = pushforward(gradient_closure, outer(backend), x, v, outer_extras)
     return p
 end
 
 function hvp_aux(f, backend, x, v, extras, ::ReverseOverForward)
     # gradient of the JVP
-    inner_extras = prepare_pushforward(extras, f, inner(backend), x)
-    jvp_closure(z) = pushforward(f, inner(backend), z, v, inner_extras)
+    function jvp_closure(z)
+        inner_extras = prepare_pushforward(extras, f, inner(backend), z)
+        return pushforward(f, inner(backend), z, v, inner_extras)
+    end
     outer_extras = prepare_gradient(extras, jvp_closure, outer(backend), x)
     p = gradient(jvp_closure, outer(backend), x, outer_extras)
     return p
 end
 
 function hvp_aux(f, backend, x, v, extras, ::ReverseOverReverse)
     # VJP of the gradient
-    inner_extras = prepare_gradient(extras, f, inner(backend), x)
-    gradient_closure(z) = gradient(f, inner(backend), z, inner_extras)
+    function gradient_closure(z)
+        inner_extras = prepare_gradient(extras, f, inner(backend), z)
+        return gradient(f, inner(backend), z, inner_extras)
+    end
     outer_extras = prepare_pullback(extras, gradient_closure, outer(backend), x)
     p = pullback(gradient_closure, outer(backend), x, v, outer_extras)
     return p
@@ -60,8 +66,10 @@ end
 function hvp_aux(f, backend, x, v, extras, ::ForwardOverForward)
     # JVPs of JVPs in theory
     # also pushforward of gradient in practice
-    inner_extras = prepare_gradient(extras, f, inner(backend), x)
-    gradient_closure(z) = gradient(f, inner(backend), z, nothing)  # TODO: fix
+    function gradient_closure(z)
+        inner_extras = prepare_gradient(extras, f, inner(backend), z)
+        return gradient(f, inner(backend), z, inner_extras)
+    end
     outer_extras = prepare_pushforward(extras, gradient_closure, outer(backend), x)
     p = pushforward(gradient_closure, outer(backend), x, v, outer_extras)
     return p
@@ -71,7 +79,7 @@ end
     hvp!!(f, p, backend, x, v, [extras]) -> p
 """
 function hvp!!(f, p, backend::AbstractADType, x, v, extras=prepare_hvp(f, backend, x))
-    new_backend = SecondOrder(backend, backend)
+    new_backend = SecondOrder(backend)
     new_extras = prepare_hvp(f, new_backend, x)
     return hvp!!(f, p, new_backend, x, v, new_extras)
 end
@@ -87,32 +95,40 @@ function hvp!!(f, p, backend::SecondOrder, x, v, extras=prepare_hvp(f, backend,
 end
 
 function hvp_aux!!(f, p, backend, x, v, extras, ::ForwardOverReverse)
-    inner_extras = prepare_gradient(extras, f, inner(backend), x)
-    gradient_closure(z) = gradient(f, inner(backend), z, inner_extras)
+    function gradient_closure(z)
+        inner_extras = prepare_gradient(extras, f, inner(backend), z)
+        return gradient(f, inner(backend), z, inner_extras)
+    end
     outer_extras = prepare_pushforward(extras, gradient_closure, outer(backend), x)
     p = pushforward!!(gradient_closure, p, outer(backend), x, v, outer_extras)
     return p
 end
 
 function hvp_aux!!(f, p, backend, x, v, extras, ::ReverseOverForward)
-    inner_extras = prepare_pushforward(extras, f, inner(backend), x)
-    jvp_closure(z) = pushforward(f, inner(backend), z, v, inner_extras)
+    function jvp_closure(z)
+        inner_extras = prepare_pushforward(extras, f, inner(backend), z)
+        return pushforward(f, inner(backend), z, v, inner_extras)
+    end
     outer_extras = prepare_gradient(extras, jvp_closure, outer(backend), x)
     p = gradient!!(jvp_closure, p, outer(backend), x, outer_extras)
     return p
 end
 
 function hvp_aux!!(f, p, backend, x, v, extras, ::ReverseOverReverse)
-    inner_extras = prepare_gradient(extras, f, inner(backend), x)
-    gradient_closure(z) = gradient(f, inner(backend), z, inner_extras)
+    function gradient_closure(z)
+        inner_extras = prepare_gradient(extras, f, inner(backend), z)
+        return gradient(f, inner(backend), z, inner_extras)
+    end
     outer_extras = prepare_pullback(extras, gradient_closure, outer(backend), x)
     p = pullback!!(gradient_closure, p, outer(backend), x, v, outer_extras)
     return p
 end
 
 function hvp_aux!!(f, p, backend, x, v, extras, ::ForwardOverForward)
-    inner_extras = prepare_gradient(extras, f, inner(backend), x)
-    gradient_closure(z) = gradient(f, inner(backend), z, nothing)  # TODO: fix
+    function gradient_closure(z)
+        inner_extras = prepare_gradient(extras, f, inner(backend), z)
+        return gradient(f, inner(backend), z, inner_extras)
+    end
     outer_extras = prepare_pushforward(extras, gradient_closure, outer(backend), x)
     p = pushforward!!(gradient_closure, p, outer(backend), x, v, outer_extras)
     return p

src/second_derivative.jl

@@ -4,18 +4,20 @@
     second_derivative(f, backend, x, [extras]) -> der2
 """
 function second_derivative(
-    f, backend::AbstractADType, x::Number, extras=prepare_second_derivative(f, backend, x)
+    f, backend::AbstractADType, x, extras=prepare_second_derivative(f, backend, x)
 )
     new_backend = SecondOrder(backend, backend)
     new_extras = prepare_second_derivative(f, new_backend, x)
     return second_derivative(f, new_backend, x, new_extras)
 end
 
 function second_derivative(
-    f, backend::SecondOrder, x::Number, extras=prepare_second_derivative(f, backend, x)
+    f, backend::SecondOrder, x, extras=prepare_second_derivative(f, backend, x)
 )
-    inner_extras = prepare_derivative(extras, f, inner(backend), x)
-    derivative_closure(z) = derivative(f, inner(backend), z, inner_extras)
+    function derivative_closure(z)
+        inner_extras = prepare_derivative(extras, f, inner(backend), z)
+        return derivative(f, inner(backend), z, inner_extras)
+    end
     outer_extras = prepare_derivative(extras, derivative_closure, outer(backend), x)
     der2 = derivative(derivative_closure, outer(backend), x, outer_extras)
     return der2
@@ -25,26 +27,20 @@ end
     second_derivative!!(f, der2, backend, x, [extras]) -> der2
 """
 function second_derivative!!(
-    f,
-    der2,
-    backend::AbstractADType,
-    x::Number,
-    extras=prepare_second_derivative(f, backend, x),
+    f, der2, backend::AbstractADType, x, extras=prepare_second_derivative(f, backend, x)
 )
     new_backend = SecondOrder(backend, backend)
     new_extras = prepare_second_derivative(f, new_backend, x)
     return second_derivative!!(f, der2, new_backend, x, new_extras)
 end
 
 function second_derivative!!(
-    f,
-    der2,
-    backend::SecondOrder,
-    x::Number,
-    extras=prepare_second_derivative(f, backend, x),
+    f, der2, backend::SecondOrder, x, extras=prepare_second_derivative(f, backend, x)
 )
-    inner_extras = prepare_derivative(extras, f, inner(backend), x)
-    derivative_closure(z) = derivative(f, inner(backend), z, inner_extras)
+    function derivative_closure(z)
+        inner_extras = prepare_derivative(extras, f, inner(backend), z)
+        return derivative(f, inner(backend), z, inner_extras)
+    end
     outer_extras = prepare_derivative(extras, derivative_closure, outer(backend), x)
     der2 = derivative!!(derivative_closure, der2, outer(backend), x, outer_extras)
     return der2

src/second_order.jl

@@ -14,6 +14,8 @@ struct SecondOrder{AD1<:AbstractADType,AD2<:AbstractADType} <: AbstractADType
     inner::AD2
 end
 
+SecondOrder(backend::AbstractADType) = SecondOrder(backend, backend)
+
 inner(backend::SecondOrder) = backend.inner
 outer(backend::SecondOrder) = backend.outer