diff --git a/DifferentiationInterface/Project.toml b/DifferentiationInterface/Project.toml index 9446f551c..b1b33ae8c 100644 --- a/DifferentiationInterface/Project.toml +++ b/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" diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceFiniteDiffExt/DifferentiationInterfaceFiniteDiffExt.jl b/DifferentiationInterface/ext/DifferentiationInterfaceFiniteDiffExt/DifferentiationInterfaceFiniteDiffExt.jl index 4970eb96e..d7441adfb 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceFiniteDiffExt/DifferentiationInterfaceFiniteDiffExt.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceFiniteDiffExt/DifferentiationInterfaceFiniteDiffExt.jl @@ -7,6 +7,7 @@ using FiniteDiff: GradientCache, HessianCache, JacobianCache, + JVPCache, finite_difference_derivative, finite_difference_gradient, finite_difference_gradient!, @@ -14,6 +15,8 @@ using FiniteDiff: finite_difference_hessian!, finite_difference_jacobian, finite_difference_jacobian!, + finite_difference_jvp, + finite_difference_jvp!, default_relstep using LinearAlgebra: dot, mul! diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceFiniteDiffExt/onearg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceFiniteDiffExt/onearg.jl index e99b70ad0..b25b6f2db 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceFiniteDiffExt/onearg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceFiniteDiffExt/onearg.jl @@ -1,6 +1,7 @@ ## Pushforward -struct FiniteDiffOneArgPushforwardPrep{R,A} <: DI.PushforwardPrep +struct FiniteDiffOneArgPushforwardPrep{C,R,A} <: DI.PushforwardPrep + cache::C relstep::R absstep::A end @@ -8,6 +9,13 @@ 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 diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceFiniteDiffExt/twoarg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceFiniteDiffExt/twoarg.jl index 2b62ec243..84bf2589e 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceFiniteDiffExt/twoarg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceFiniteDiffExt/twoarg.jl @@ -1,6 +1,7 @@ ## Pushforward -struct FiniteDiffTwoArgPushforwardPrep{R,A} <: DI.PushforwardPrep +struct FiniteDiffTwoArgPushforwardPrep{C,R,A} <: DI.PushforwardPrep + cache::C relstep::R absstep::A end @@ -8,6 +9,11 @@ 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 @@ function DI.prepare_pushforward( 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 @@ function DI.value_and_pushforward( 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 diff --git a/DifferentiationInterface/test/Back/FiniteDiff/benchmark.jl b/DifferentiationInterface/test/Back/FiniteDiff/benchmark.jl new file mode 100644 index 000000000..f2783696c --- /dev/null +++ b/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 diff --git a/DifferentiationInterface/test/Back/FiniteDiff/test.jl b/DifferentiationInterface/test/Back/FiniteDiff/test.jl index 5276d6233..34a8c80e5 100644 --- a/DifferentiationInterface/test/Back/FiniteDiff/test.jl +++ b/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) diff --git a/DifferentiationInterface/test/Back/ForwardDiff/benchmark.jl b/DifferentiationInterface/test/Back/ForwardDiff/benchmark.jl new file mode 100644 index 000000000..aea81d03b --- /dev/null +++ b/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 diff --git a/DifferentiationInterface/test/Back/ForwardDiff/test.jl b/DifferentiationInterface/test/Back/ForwardDiff/test.jl index 061836a81..f4024afbd 100644 --- a/DifferentiationInterface/test/Back/ForwardDiff/test.jl +++ b/DifferentiationInterface/test/Back/ForwardDiff/test.jl @@ -28,55 +28,49 @@ for backend in backends @test check_inplace(backend) end -## Dense +@testset "Dense" begin + test_differentiation( + backends, default_scenarios(; include_constantified=true); logging=LOGGING + ) -test_differentiation( - backends, default_scenarios(; include_constantified=true); logging=LOGGING -); + test_differentiation( + AutoForwardDiff(), + default_scenarios(; + include_normal=false, include_batchified=false, include_cachified=true + ); + logging=LOGGING, + ) + + test_differentiation( + AutoForwardDiff(); correctness=false, type_stability=:prepared, logging=LOGGING + ) + + test_differentiation( + AutoForwardDiff(; chunksize=5); + correctness=false, + type_stability=:full, + excluded=[:hessian], + logging=LOGGING, + ) +end + +@testset "Sparse" begin + test_differentiation( + MyAutoSparse(AutoForwardDiff()), default_scenarios(); logging=LOGGING + ) + + test_differentiation( + MyAutoSparse(AutoForwardDiff()), + sparse_scenarios(; include_constantified=true); + sparsity=true, + logging=LOGGING, + ) +end + +@testset "Weird" begin + test_differentiation(AutoForwardDiff(), component_scenarios(); logging=LOGGING) + test_differentiation(AutoForwardDiff(), static_scenarios(); logging=LOGGING) -test_differentiation( - AutoForwardDiff(), - default_scenarios(; - include_normal=false, include_batchified=false, include_cachified=true - ); - logging=LOGGING, -); - -test_differentiation( - AutoForwardDiff(); correctness=false, type_stability=:prepared, logging=LOGGING -); - -test_differentiation( - AutoForwardDiff(; chunksize=5); - correctness=false, - type_stability=:full, - excluded=[:hessian], - logging=LOGGING, -); - -test_differentiation( - backends, - vcat(component_scenarios(), static_scenarios()); # FD accesses individual indices - excluded=vcat(SECOND_ORDER, [:jacobian]), # jacobian is super slow for some reason - logging=LOGGING, -); - -## Sparse - -test_differentiation(MyAutoSparse(AutoForwardDiff()), default_scenarios(); logging=LOGGING); - -test_differentiation( - MyAutoSparse(AutoForwardDiff()), - sparse_scenarios(; include_constantified=true); - sparsity=true, - logging=LOGGING, -); - -## Static - -test_differentiation(AutoForwardDiff(), static_scenarios(); logging=LOGGING) - -@testset verbose = true "StaticArrays" begin @testset "Batch size" begin @test DI.pick_batchsize(AutoForwardDiff(), rand(7)) isa DI.BatchSizeSettings{7} @test DI.pick_batchsize(AutoForwardDiff(; chunksize=5), rand(7)) isa @@ -87,23 +81,7 @@ test_differentiation(AutoForwardDiff(), static_scenarios(); logging=LOGGING) AutoForwardDiff(; chunksize=5), @SVector(rand(7)) )) isa DI.BatchSizeSettings{5} end - - 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; +end @testset verbose = true "Overloaded inputs" begin backend = AutoForwardDiff()