In-place second order stuff (#117)

Author: gdalle

README.md

@@ -11,8 +11,8 @@ An interface to various automatic differentiation backends in Julia.
 
 This package provides a backend-agnostic syntax to differentiate functions of the following types:
 
-- **allocating**: `f(x) = y`
-- **mutating**: `f!(y, x) = nothing`
+- _allocating_: `f(x) = y`
+- _mutating_: `f!(y, x) = nothing`
 
 ## Features
 

docs/src/backends.md

@@ -67,10 +67,10 @@ AutoFastDifferentiation
 You can use [`check_available`](@ref) to verify whether a given backend is loaded, like we did below:
 
 ```@example backends
-header = "| Backend | available |"  # hide
+header = "| backend | available |"  # hide
 subheader = "|---|---|"  # hide
 rows = map(all_backends()) do backend  # hide
-    "| `$(backend_string(backend))` | $(check_available(backend) ? '✓' : '✗') |"  # hide
+    "| `$(backend_string(backend))` | $(check_available(backend) ? '✅' : '❌') |"  # hide
 end  # hide
 Markdown.parse(join(vcat(header, subheader, rows...), "\n"))  # hide
 ```
@@ -82,10 +82,10 @@ Only some are compatible with mutating functions `f!(y, x) = nothing`.
 You can use [`check_mutation`](@ref) to check that feature, like we did below:
 
 ```@example backends
-header = "| Backend | mutation |"  # hide
+header = "| backend | mutation |"  # hide
 subheader = "|---|---|"  # hide
 rows = map(all_backends()) do backend  # hide
-    "| `$(backend_string(backend))` | $(check_mutation(backend) ? '✓' : '✗') |"  # hide
+    "| `$(backend_string(backend))` | $(check_mutation(backend) ? '✅' : '❌') |"  # hide
 end  # hide
 Markdown.parse(join(vcat(header, subheader, rows...), "\n"))  # hide
 ```

docs/src/overview.md

@@ -56,6 +56,14 @@ Second-order differentiation is also supported, with the following operators:
 | [`hvp`](@ref)               | `Any`           | `Number`     | same as `x`      | `size(x)`                |
 | [`hessian`](@ref)           | `AbstractArray` | `Number`     | `AbstractMatrix` | `(length(x), length(x))` |
 
+We only define two variants for now:
+
+| out-of-place                | in-place (or not)             |
+| --------------------------- | ----------------------------- |
+| [`second_derivative`](@ref) | [`second_derivative!!`](@ref) |
+| [`hvp`](@ref)               | [`hvp!!`](@ref)               |
+| [`hessian`](@ref)           | [`hessian!!`](@ref)           |
+
 !!! danger
     This is an experimental functionality, use at your own risk.
 

ext/DifferentiationInterfaceEnzymeExt/reverse_allocating.jl

@@ -56,5 +56,8 @@ function DI.gradient(f, ::AutoReverseEnzyme, x::AbstractArray, extras::Nothing)
 end
 
 function DI.gradient!!(f, grad, ::AutoReverseEnzyme, x::AbstractArray, extras::Nothing)
-    return gradient!(Reverse, grad, f, x)
+    grad_sametype = convert(typeof(x), grad)
+    gradient!(Reverse, grad_sametype, f, x)
+    grad .= grad_sametype
+    return grad
 end

lib/DifferentiationInterfaceTest/src/tests/benchmark.jl

@@ -306,7 +306,7 @@ function run_benchmark!(
 )
     (; f, x, y, dy) = deepcopy(scen)
     extras = prepare_second_derivative(f, ba, x)
-    bench1 = @be second_derivative(f, ba, x, extras)
+    bench1 = @be mysimilar(dy) second_derivative!!(f, _, ba, x, extras)
     # only test allocations if the output is scalar
     if allocations && y isa Number
         @test 0 == minimum(bench1).allocs
@@ -326,7 +326,7 @@ function run_benchmark!(
 )
     (; f, x, y, dx) = deepcopy(scen)
     extras = prepare_hvp(f, ba, x)
-    bench1 = @be hvp(f, ba, x, dx, extras)
+    bench1 = @be mysimilar(dx) hvp!!(f, _, ba, x, dx, extras)
     # no test for now
     record!(data, ba, op, hvp, scen, bench1)
     return nothing
@@ -343,7 +343,8 @@ function run_benchmark!(
 )
     (; f, x, y) = deepcopy(scen)
     extras = prepare_hessian(f, ba, x)
-    bench1 = @be hessian(f, ba, x, extras)
+    hess_template = Matrix{typeof(y)}(undef, length(x), length(x))
+    bench1 = @be similar(hess_template) hessian!!(f, _, ba, x, extras)
     # no test for now
     record!(data, ba, op, hessian, scen, bench1)
     return nothing

lib/DifferentiationInterfaceTest/src/tests/call_count.jl

@@ -138,10 +138,11 @@ end
 function test_call_count(
     ba::AbstractADType, ::typeof(second_derivative), scen::Scenario{false}
 )
-    (; f, x, y) = deepcopy(scen)
+    (; f, x, y, dy) = deepcopy(scen)
     extras = prepare_second_derivative(CallCounter(f), ba, x)
     cc = CallCounter(f)
-    second_derivative(cc, ba, x, extras)
+    der2_in = mysimilar(dy)
+    second_derivative!!(cc, der2_in, ba, x, extras)
     # what to test?
     return nothing
 end
@@ -152,7 +153,8 @@ function test_call_count(ba::AbstractADType, ::typeof(hvp), scen::Scenario{false
     (; f, x, y, dx) = deepcopy(scen)
     extras = prepare_hvp(CallCounter(f), ba, x)
     cc = CallCounter(f)
-    hvp(cc, ba, x, dx, extras)
+    p_in = mysimilar(dx)
+    hvp!!(cc, p_in, ba, x, dx, extras)
     # what to test?
     return nothing
 end
@@ -163,7 +165,8 @@ function test_call_count(ba::AbstractADType, ::typeof(hessian), scen::Scenario{f
     (; f, x, y, dx) = deepcopy(scen)
     extras = prepare_hessian(CallCounter(f), ba, x)
     cc = CallCounter(f)
-    hessian(cc, ba, x, extras)
+    hess_in = Matrix{typeof(y)}(undef, length(x), length(x))
+    hessian!!(cc, hess_in, ba, x, extras)
     # what to test?
     return nothing
 end

lib/DifferentiationInterfaceTest/src/tests/correctness.jl

@@ -16,6 +16,7 @@ function test_correctness(
     ::typeof(pushforward),
     scen::Scenario{false};
     isapprox::Function,
+    atol,
     rtol,
 )
     (; f, x, y, dx, dy, ref) = new_scen = deepcopy(scen)
@@ -31,7 +32,7 @@ function test_correctness(
     dy3 = pushforward(f, ba, x, dx)
     dy4 = pushforward!!(f, mysimilar(dy), ba, x, dx)
 
-    let (≈)(x, y) = isapprox(x, y; rtol)
+    let (≈)(x, y) = isapprox(x, y; atol, rtol)
         @testset "Primal value" begin
             @test y1 ≈ y
             @test y2 ≈ y
@@ -52,6 +53,7 @@ function test_correctness(
     ::typeof(pushforward),
     scen::Scenario{true};
     isapprox::Function,
+    atol,
     rtol,
 )
     (; f, x, y, dx, dy, ref) = new_scen = deepcopy(scen)
@@ -65,7 +67,7 @@ function test_correctness(
     y10 = mysimilar(y)
     y1, dy1 = value_and_pushforward!!(f!, y10, mysimilar(dy), ba, x, dx)
 
-    let (≈)(x, y) = isapprox(x, y; rtol)
+    let (≈)(x, y) = isapprox(x, y; atol, rtol)
         @testset "Primal value" begin
             @test y10 ≈ y
             @test y1 ≈ y
@@ -81,7 +83,12 @@ end
 ## Pullback
 
 function test_correctness(
-    ba::AbstractADType, ::typeof(pullback), scen::Scenario{false}; isapprox::Function, rtol
+    ba::AbstractADType,
+    ::typeof(pullback),
+    scen::Scenario{false};
+    isapprox::Function,
+    atol,
+    rtol,
 )
     (; f, x, y, dx, dy, ref) = new_scen = deepcopy(scen)
     dx_true = if ref isa AbstractADType
@@ -96,7 +103,7 @@ function test_correctness(
     dx3 = pullback(f, ba, x, dy)
     dx4 = pullback!!(f, mysimilar(dx), ba, x, dy)
 
-    let (≈)(x, y) = isapprox(x, y; rtol)
+    let (≈)(x, y) = isapprox(x, y; atol, rtol)
         @testset "Primal value" begin
             @test y1 ≈ y
             @test y2 ≈ y
@@ -113,7 +120,12 @@ function test_correctness(
 end
 
 function test_correctness(
-    ba::AbstractADType, ::typeof(pullback), scen::Scenario{true}; isapprox::Function, rtol
+    ba::AbstractADType,
+    ::typeof(pullback),
+    scen::Scenario{true};
+    isapprox::Function,
+    atol,
+    rtol,
 )
     (; f, x, y, dx, dy, ref) = new_scen = deepcopy(scen)
     f! = f
@@ -126,7 +138,7 @@ function test_correctness(
     y10 = mysimilar(y)
     y1, dx1 = value_and_pullback!!(f!, y10, mysimilar(dx), ba, x, dy)
 
-    let (≈)(x, y) = isapprox(x, y; rtol)
+    let (≈)(x, y) = isapprox(x, y; atol, rtol)
         @testset "Primal value" begin
             @test y10 ≈ y
             @test y1 ≈ y
@@ -146,6 +158,7 @@ function test_correctness(
     ::typeof(derivative),
     scen::Scenario{false};
     isapprox::Function,
+    atol,
     rtol,
 )
     (; f, x, y, dx, dy, ref) = new_scen = deepcopy(scen)
@@ -161,7 +174,7 @@ function test_correctness(
     der3 = derivative(f, ba, x)
     der4 = derivative!!(f, mysimilar(dy), ba, x)
 
-    let (≈)(x, y) = isapprox(x, y; rtol)
+    let (≈)(x, y) = isapprox(x, y; atol, rtol)
         @testset "Primal value" begin
             @test y1 ≈ y
             @test y2 ≈ y
@@ -178,7 +191,12 @@ function test_correctness(
 end
 
 function test_correctness(
-    ba::AbstractADType, ::typeof(derivative), scen::Scenario{true}; isapprox::Function, rtol
+    ba::AbstractADType,
+    ::typeof(derivative),
+    scen::Scenario{true};
+    isapprox::Function,
+    atol,
+    rtol,
 )
     (; f, x, y, dx, dy, ref) = new_scen = deepcopy(scen)
     f! = f
@@ -191,7 +209,7 @@ function test_correctness(
     y10 = mysimilar(y)
     y1, der1 = value_and_derivative!!(f!, y10, mysimilar(dy), ba, x)
 
-    let (≈)(x, y) = isapprox(x, y; rtol)
+    let (≈)(x, y) = isapprox(x, y; atol, rtol)
         @testset "Primal value" begin
             @test y10 ≈ y
             @test y1 ≈ y
@@ -207,7 +225,12 @@ end
 ## Gradient
 
 function test_correctness(
-    ba::AbstractADType, ::typeof(gradient), scen::Scenario{false}; isapprox::Function, rtol
+    ba::AbstractADType,
+    ::typeof(gradient),
+    scen::Scenario{false};
+    isapprox::Function,
+    atol,
+    rtol,
 )
     (; f, x, y, dx, dy, ref) = new_scen = deepcopy(scen)
     grad_true = if ref isa AbstractADType
@@ -222,7 +245,7 @@ function test_correctness(
     grad3 = gradient(f, ba, x)
     grad4 = gradient!!(f, mysimilar(dx), ba, x)
 
-    let (≈)(x, y) = isapprox(x, y; rtol)
+    let (≈)(x, y) = isapprox(x, y; atol, rtol)
         @testset "Primal value" begin
             @test y1 ≈ y
             @test y2 ≈ y
@@ -241,7 +264,12 @@ end
 ## Jacobian
 
 function test_correctness(
-    ba::AbstractADType, ::typeof(jacobian), scen::Scenario{false}; isapprox::Function, rtol
+    ba::AbstractADType,
+    ::typeof(jacobian),
+    scen::Scenario{false};
+    isapprox::Function,
+    atol,
+    rtol,
 )
     (; f, x, y, ref) = new_scen = deepcopy(scen)
     jac_true = if ref isa AbstractADType
@@ -256,7 +284,7 @@ function test_correctness(
     jac3 = jacobian(f, ba, x)
     jac4 = jacobian!!(f, mysimilar(jac_true), ba, x)
 
-    let (≈)(x, y) = isapprox(x, y; rtol)
+    let (≈)(x, y) = isapprox(x, y; atol, rtol)
         @testset "Primal value" begin
             @test y1 ≈ y
             @test y2 ≈ y
@@ -273,7 +301,12 @@ function test_correctness(
 end
 
 function test_correctness(
-    ba::AbstractADType, ::typeof(jacobian), scen::Scenario{true}; isapprox::Function, rtol
+    ba::AbstractADType,
+    ::typeof(jacobian),
+    scen::Scenario{true};
+    isapprox::Function,
+    atol,
+    rtol,
 )
     (; f, x, y, dy, ref) = new_scen = deepcopy(scen)
     f! = f
@@ -287,7 +320,7 @@ function test_correctness(
     y10 = mysimilar(y)
     y1, jac1 = value_and_jacobian!!(f!, y10, mysimilar(jac_true), ba, x)
 
-    let (≈)(x, y) = isapprox(x, y; rtol)
+    let (≈)(x, y) = isapprox(x, y; atol, rtol)
         @testset "Primal value" begin
             @test y10 ≈ y
             @test y1 ≈ y
@@ -307,20 +340,23 @@ function test_correctness(
     ::typeof(second_derivative),
     scen::Scenario;
     isapprox::Function,
+    atol,
     rtol,
 )
-    (; f, x, ref) = new_scen = deepcopy(scen)
+    (; f, x, dy, ref) = new_scen = deepcopy(scen)
     der2_true = if ref isa AbstractADType
         second_derivative(f, ref, x)
     else
         ref.second_derivative(x)
     end
 
     der21 = second_derivative(f, ba, x)
+    der22 = second_derivative!!(f, mysimilar(dy), ba, x)
 
-    let (≈)(x, y) = isapprox(x, y; rtol)
+    let (≈)(x, y) = isapprox(x, y; atol, rtol)
         @testset "Second derivative value" begin
             @test der21 ≈ der2_true
+            @test der22 ≈ der2_true
         end
     end
     test_scen_intact(new_scen, scen)
@@ -330,7 +366,7 @@ end
 ## Hessian-vector product
 
 function test_correctness(
-    ba::AbstractADType, ::typeof(hvp), scen::Scenario; isapprox::Function, rtol
+    ba::AbstractADType, ::typeof(hvp), scen::Scenario; isapprox::Function, atol, rtol
 )
     (; f, x, dx, ref) = new_scen = deepcopy(scen)
     hvp_true = if ref isa AbstractADType
@@ -340,10 +376,12 @@ function test_correctness(
     end
 
     hvp1 = hvp(f, ba, x, dx)
+    hvp2 = hvp!!(f, mysimilar(dx), ba, x, dx)
 
-    let (≈)(x, y) = isapprox(x, y; rtol)
+    let (≈)(x, y) = isapprox(x, y; atol, rtol)
         @testset "HVP value" begin
             @test hvp1 ≈ hvp_true
+            @test hvp2 ≈ hvp_true
         end
     end
     test_scen_intact(new_scen, scen)
@@ -353,7 +391,7 @@ end
 ## Hessian
 
 function test_correctness(
-    ba::AbstractADType, ::typeof(hessian), scen::Scenario; isapprox::Function, rtol
+    ba::AbstractADType, ::typeof(hessian), scen::Scenario; isapprox::Function, atol, rtol
 )
     (; f, x, y, ref) = new_scen = deepcopy(scen)
     hess_true = if ref isa AbstractADType
@@ -363,10 +401,12 @@ function test_correctness(
     end
 
     hess1 = hessian(f, ba, x)
+    hess2 = hessian!!(f, mysimilar(hess_true), ba, x)
 
-    let (≈)(x, y) = isapprox(x, y; rtol)
+    let (≈)(x, y) = isapprox(x, y; atol, rtol)
         @testset "Hessian value" begin
             @test hess1 ≈ hess_true
+            @test hess2 ≈ hess_true
         end
     end
     test_scen_intact(new_scen, scen)