feat: use JVPCache for FiniteDiff pushforwards

DifferentiationInterface/Project.toml

@@ -1,7 +1,7 @@
 name = "DifferentiationInterface"
 uuid = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 authors = ["Guillaume Dalle", "Adrian Hill"]
-version = "0.6.37"
+version = "0.6.38"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -56,7 +56,7 @@ Enzyme = "0.13.17"
 EnzymeCore = "0.8.8"
 ExplicitImports = "1.10.1"
 FastDifferentiation = "0.4.3"
-FiniteDiff = "2.23.1"
+FiniteDiff = "2.27.0"
 FiniteDifferences = "0.12.31"
 ForwardDiff = "0.10.36"
 GTPSA = "1.4.0"

DifferentiationInterface/ext/DifferentiationInterfaceFiniteDiffExt/DifferentiationInterfaceFiniteDiffExt.jl

@@ -7,13 +7,16 @@ using FiniteDiff:
     GradientCache,
     HessianCache,
     JacobianCache,
+    JVPCache,
     finite_difference_derivative,
     finite_difference_gradient,
     finite_difference_gradient!,
     finite_difference_hessian,
     finite_difference_hessian!,
     finite_difference_jacobian,
     finite_difference_jacobian!,
+    finite_difference_jvp,
+    finite_difference_jvp!,
     default_relstep
 using LinearAlgebra: dot, mul!
 

DifferentiationInterface/ext/DifferentiationInterfaceFiniteDiffExt/onearg.jl

@@ -1,13 +1,21 @@
 ## Pushforward
 
-struct FiniteDiffOneArgPushforwardPrep{R,A} <: DI.PushforwardPrep
+struct FiniteDiffOneArgPushforwardPrep{C,R,A} <: DI.PushforwardPrep
+    cache::C
     relstep::R
     absstep::A
 end
 
 function DI.prepare_pushforward(
     f, backend::AutoFiniteDiff, x, tx::NTuple, contexts::Vararg{DI.Context,C}
 ) where {C}
+    fc = DI.with_contexts(f, contexts...)
+    y = fc(x)
+    cache = if x isa Number || y isa Number
+        nothing
+    else
+        JVPCache(similar(x), y, fdtype(backend))
+    end
     relstep = if isnothing(backend.relstep)
         default_relstep(fdtype(backend), eltype(x))
     else
@@ -18,12 +26,12 @@ function DI.prepare_pushforward(
     else
         backend.relstep
     end
-    return FiniteDiffOneArgPushforwardPrep(relstep, absstep)
+    return FiniteDiffOneArgPushforwardPrep(cache, relstep, absstep)
 end
 
 function DI.pushforward(
     f,
-    prep::FiniteDiffOneArgPushforwardPrep,
+    prep::FiniteDiffOneArgPushforwardPrep{Nothing},
     backend::AutoFiniteDiff,
     x,
     tx::NTuple,
@@ -41,7 +49,7 @@ end
 
 function DI.value_and_pushforward(
     f,
-    prep::FiniteDiffOneArgPushforwardPrep,
+    prep::FiniteDiffOneArgPushforwardPrep{Nothing},
     backend::AutoFiniteDiff,
     x,
     tx::NTuple,
@@ -64,6 +72,39 @@ function DI.value_and_pushforward(
     return y, ty
 end
 
+function DI.pushforward(
+    f,
+    prep::FiniteDiffOneArgPushforwardPrep{<:JVPCache},
+    ::AutoFiniteDiff,
+    x,
+    tx::NTuple,
+    contexts::Vararg{DI.Context,C},
+) where {C}
+    (; relstep, absstep) = prep
+    fc = DI.with_contexts(f, contexts...)
+    ty = map(tx) do dx
+        finite_difference_jvp(fc, x, dx, prep.cache; relstep, absstep)
+    end
+    return ty
+end
+
+function DI.value_and_pushforward(
+    f,
+    prep::FiniteDiffOneArgPushforwardPrep{<:JVPCache},
+    ::AutoFiniteDiff,
+    x,
+    tx::NTuple,
+    contexts::Vararg{DI.Context,C},
+) where {C}
+    (; relstep, absstep) = prep
+    fc = DI.with_contexts(f, contexts...)
+    y = fc(x)
+    ty = map(tx) do dx
+        finite_difference_jvp(fc, x, dx, prep.cache, y; relstep, absstep)
+    end
+    return y, ty
+end
+
 ## Derivative
 
 struct FiniteDiffOneArgDerivativePrep{C,R,A} <: DI.DerivativePrep

DifferentiationInterface/ext/DifferentiationInterfaceFiniteDiffExt/twoarg.jl

@@ -1,13 +1,19 @@
 ## Pushforward
 
-struct FiniteDiffTwoArgPushforwardPrep{R,A} <: DI.PushforwardPrep
+struct FiniteDiffTwoArgPushforwardPrep{C,R,A} <: DI.PushforwardPrep
+    cache::C
     relstep::R
     absstep::A
 end
 
 function DI.prepare_pushforward(
     f!, y, backend::AutoFiniteDiff, x, tx::NTuple, contexts::Vararg{DI.Context,C}
 ) where {C}
+    cache = if x isa Number
+        nothing
+    else
+        JVPCache(similar(x), similar(y), fdtype(backend))
+    end
     relstep = if isnothing(backend.relstep)
         default_relstep(fdtype(backend), eltype(x))
     else
@@ -18,14 +24,13 @@
     else
         backend.relstep
     end
-    return FiniteDiffTwoArgPushforwardPrep(relstep, absstep)
-    return DI.NoPushforwardPrep()
+    return FiniteDiffTwoArgPushforwardPrep(cache, relstep, absstep)
 end
 
 function DI.value_and_pushforward(
     f!,
     y,
-    prep::FiniteDiffTwoArgPushforwardPrep,
+    prep::FiniteDiffTwoArgPushforwardPrep{Nothing},
     backend::AutoFiniteDiff,
     x,
     tx::NTuple,
@@ -52,6 +57,84 @@
     return y, ty
 end
 
+function DI.pushforward(
+    f!,
+    y,
+    prep::FiniteDiffTwoArgPushforwardPrep{<:JVPCache},
+    ::AutoFiniteDiff,
+    x,
+    tx::NTuple,
+    contexts::Vararg{DI.Context,C},
+) where {C}
+    (; relstep, absstep) = prep
+    fc! = DI.with_contexts(f!, contexts...)
+    ty = map(tx) do dx
+        dy = similar(y)
+        finite_difference_jvp!(dy, fc!, x, dx, prep.cache; relstep, absstep)
+        dy
+    end
+    return ty
+end
+
+function DI.value_and_pushforward(
+    f!,
+    y,
+    prep::FiniteDiffTwoArgPushforwardPrep{<:JVPCache},
+    ::AutoFiniteDiff,
+    x,
+    tx::NTuple,
+    contexts::Vararg{DI.Context,C},
+) where {C}
+    (; relstep, absstep) = prep
+    fc! = DI.with_contexts(f!, contexts...)
+    ty = map(tx) do dx
+        dy = similar(y)
+        finite_difference_jvp!(dy, fc!, x, dx, prep.cache; relstep, absstep)
+        dy
+    end
+    fc!(y, x)
+    return y, ty
+end
+
+function DI.pushforward!(
+    f!,
+    y,
+    ty::NTuple,
+    prep::FiniteDiffTwoArgPushforwardPrep{<:JVPCache},
+    ::AutoFiniteDiff,
+    x,
+    tx::NTuple,
+    contexts::Vararg{DI.Context,C},
+) where {C}
+    (; relstep, absstep) = prep
+    fc! = DI.with_contexts(f!, contexts...)
+    for b in eachindex(tx, ty)
+        dx, dy = tx[b], ty[b]
+        finite_difference_jvp!(dy, fc!, x, dx, prep.cache; relstep, absstep)
+    end
+    return ty
+end
+
+function DI.value_and_pushforward!(
+    f!,
+    y,
+    ty::NTuple,
+    prep::FiniteDiffTwoArgPushforwardPrep{<:JVPCache},
+    ::AutoFiniteDiff,
+    x,
+    tx::NTuple,
+    contexts::Vararg{DI.Context,C},
+) where {C}
+    (; relstep, absstep) = prep
+    fc! = DI.with_contexts(f!, contexts...)
+    for b in eachindex(tx, ty)
+        dx, dy = tx[b], ty[b]
+        finite_difference_jvp!(dy, fc!, x, dx, prep.cache; relstep, absstep)
+    end
+    fc!(y, x)
+    return y, ty
+end
+
 ## Derivative
 
 struct FiniteDiffTwoArgDerivativePrep{C,R,A} <: DI.DerivativePrep

DifferentiationInterface/test/Back/FiniteDiff/benchmark.jl

@@ -0,0 +1,33 @@
+using Pkg
+Pkg.add("FiniteDiff")
+
+using ADTypes: ADTypes
+using DifferentiationInterface, DifferentiationInterfaceTest
+import DifferentiationInterface as DI
+import DifferentiationInterfaceTest as DIT
+using FiniteDiff: FiniteDiff
+using Test
+
+LOGGING = get(ENV, "CI", "false") == "false"
+
+@testset "Benchmarking sparse" begin
+    filtered_sparse_scenarios = filter(sparse_scenarios(; band_sizes=[])) do scen
+        DIT.function_place(scen) == :in &&
+            DIT.operator_place(scen) == :in &&
+            scen.x isa AbstractVector &&
+            scen.y isa AbstractVector
+    end
+
+    data = benchmark_differentiation(
+        MyAutoSparse(AutoFiniteDiff()),
+        filtered_sparse_scenarios;
+        benchmark=:prepared,
+        excluded=SECOND_ORDER,
+        logging=LOGGING,
+    )
+    @testset "Analyzing benchmark results" begin
+        @testset "$(row[:scenario])" for row in eachrow(data)
+            @test row[:allocs] == 0
+        end
+    end
+end

DifferentiationInterface/test/Back/FiniteDiff/test.jl

@@ -17,21 +17,32 @@ for backend in [AutoFiniteDiff()]
     @test check_inplace(backend)
 end
 
-test_differentiation(
-    AutoFiniteDiff(),
-    default_scenarios(; include_constantified=true, include_cachified=true);
-    excluded=[:second_derivative, :hvp],
-    logging=LOGGING,
-);
-
-test_differentiation(
-    [
-        AutoFiniteDiff(; relstep=cbrt(eps(Float64))),
-        AutoFiniteDiff(; relstep=cbrt(eps(Float64)), absstep=cbrt(eps(Float64))),
-    ];
-    excluded=[:second_derivative, :hvp],
-    logging=LOGGING,
-);
+@testset "Dense" begin
+    test_differentiation(
+        AutoFiniteDiff(),
+        default_scenarios(; include_constantified=true, include_cachified=true);
+        excluded=[:second_derivative, :hvp],
+        logging=LOGGING,
+    )
+
+    test_differentiation(
+        [
+            AutoFiniteDiff(; relstep=cbrt(eps(Float64))),
+            AutoFiniteDiff(; relstep=cbrt(eps(Float64)), absstep=cbrt(eps(Float64))),
+        ];
+        excluded=[:second_derivative, :hvp],
+        logging=LOGGING,
+    )
+end
+
+@testset "Sparse" begin
+    test_differentiation(
+        MyAutoSparse(AutoFiniteDiff()),
+        sparse_scenarios();
+        excluded=SECOND_ORDER,
+        logging=LOGGING,
+    )
+end
 
 @testset "Complex" begin
     test_differentiation(AutoFiniteDiff(), complex_scenarios(); logging=LOGGING)

DifferentiationInterface/test/Back/ForwardDiff/benchmark.jl

@@ -0,0 +1,52 @@
+using Pkg
+Pkg.add("ForwardDiff")
+
+using ADTypes: ADTypes
+using DifferentiationInterface, DifferentiationInterfaceTest
+import DifferentiationInterface as DI
+import DifferentiationInterfaceTest as DIT
+using ForwardDiff: ForwardDiff
+using StaticArrays: StaticArrays, @SVector
+using Test
+
+LOGGING = get(ENV, "CI", "false") == "false"
+
+@testset verbose = true "Benchmarking static" begin
+    filtered_static_scenarios = filter(static_scenarios(; include_batchified=false)) do scen
+        DIT.function_place(scen) == :out && DIT.operator_place(scen) == :out
+    end
+    data = benchmark_differentiation(
+        AutoForwardDiff(),
+        filtered_static_scenarios;
+        benchmark=:prepared,
+        excluded=[:hessian, :pullback],  # TODO: figure this out
+        logging=LOGGING,
+    )
+    @testset "Analyzing benchmark results" begin
+        @testset "$(row[:scenario])" for row in eachrow(data)
+            @test row[:allocs] == 0
+        end
+    end
+end
+
+@testset "Benchmarking sparse" begin
+    filtered_sparse_scenarios = filter(sparse_scenarios(; band_sizes=[])) do scen
+        DIT.function_place(scen) == :in &&
+            DIT.operator_place(scen) == :in &&
+            scen.x isa AbstractVector &&
+            scen.y isa AbstractVector
+    end
+
+    data = benchmark_differentiation(
+        MyAutoSparse(AutoForwardDiff()),
+        filtered_sparse_scenarios;
+        benchmark=:prepared,
+        excluded=SECOND_ORDER,
+        logging=LOGGING,
+    )
+    @testset "Analyzing benchmark results" begin
+        @testset "$(row[:scenario])" for row in eachrow(data)
+            @test row[:allocs] == 0
+        end
+    end
+end