diff --git a/Project.toml b/Project.toml index 454827893..94384c3b9 100644 --- a/Project.toml +++ b/Project.toml @@ -37,7 +37,6 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" DifferentiationInterfaceChainRulesCoreExt = "ChainRulesCore" DifferentiationInterfaceChairmarksExt = ["Chairmarks"] DifferentiationInterfaceComponentArraysExt = ["ComponentArrays"] -DifferentiationInterfaceCorrectnessTestExt = ["ForwardDiff", "Zygote"] DifferentiationInterfaceDiffractorExt = ["AbstractDifferentiation", "Diffractor"] DifferentiationInterfaceEnzymeExt = "Enzyme" DifferentiationInterfaceFastDifferentiationExt = ["FastDifferentiation", "RuntimeGeneratedFunctions"] diff --git a/docs/src/tutorial.md b/docs/src/tutorial.md index c1389e2bd..537d68aec 100644 --- a/docs/src/tutorial.md +++ b/docs/src/tutorial.md @@ -129,25 +129,31 @@ It's blazingly fast. And you know what's even better? You didn't need to look at the docs of either ForwardDiff.jl or Enzyme.jl to achieve top performance with both, or to compare them. -## Benchmarking +## Testing and benchmarking -DifferentiationInterface.jl also provides some benchmarking utilities, for more involved comparison between backends. +DifferentiationInterface.jl also provides some utilities for more involved comparison between backends. They are gathered in a submodule. ```@repl tuto using DifferentiationInterface.DifferentiationTest ``` -We can now use [`test_differentiation`](@ref) in benchmarking mode, and convert its output to a `DataFrame` from [DataFrames.jl](https://github.com/JuliaData/DataFrames.jl) (disregard the test logging): +The main entry point is [`test_differentiation`](@ref), which is used as follows: ```@repl tuto data = test_differentiation( [AutoForwardDiff(), AutoEnzyme(Enzyme.Reverse)], # backends to compare [gradient], # operators to try [Scenario(f; x=x)]; # test scenario - correctness=false, # we don't care about checking the result - benchmark=true, # we care about measuring + correctness=AutoZygote(), # compare results to a "ground truth" from Zygote + benchmark=true, # measure runtime and allocations too + detailed=true, # print detailed test set ); +``` + +The output of `test_differentiation` when `benchmark=true` can be converted to a `DataFrame` from [DataFrames.jl](https://github.com/JuliaData/DataFrames.jl): + +```@repl tuto df = DataFrames.DataFrame(pairs(data)...) ``` diff --git a/ext/DifferentiationInterfaceChairmarksExt/DifferentiationInterfaceChairmarksExt.jl b/ext/DifferentiationInterfaceChairmarksExt/DifferentiationInterfaceChairmarksExt.jl index 01ea408d4..fc9e291b0 100644 --- a/ext/DifferentiationInterfaceChairmarksExt/DifferentiationInterfaceChairmarksExt.jl +++ b/ext/DifferentiationInterfaceChairmarksExt/DifferentiationInterfaceChairmarksExt.jl @@ -155,7 +155,7 @@ function run_benchmark!( ) (; f, x, y) = deepcopy(scen) extras = prepare_jacobian(f, ba, x) - jac_template = similar(y, length(y), length(x)) + jac_template = Matrix{eltype(y)}(undef, length(y), length(x)) bench1 = @be mysimilar(jac_template) value_and_jacobian!!(f, _, ba, x, extras) # never test allocations record!(data, ba, op, value_and_jacobian!!, scen, bench1) @@ -172,7 +172,7 @@ function run_benchmark!( (; f, x, y) = deepcopy(scen) f! = f extras = prepare_jacobian(f!, ba, y, x) - jac_template = similar(y, length(y), length(x)) + jac_template = Matrix{eltype(y)}(undef, length(y), length(x)) bench1 = @be (mysimilar(y), mysimilar(jac_template)) value_and_jacobian!!( f!, _[1], _[2], ba, x, extras ) diff --git a/ext/DifferentiationInterfaceComponentArraysExt/DifferentiationInterfaceComponentArraysExt.jl b/ext/DifferentiationInterfaceComponentArraysExt/DifferentiationInterfaceComponentArraysExt.jl index b6cead8ed..eacec343d 100644 --- a/ext/DifferentiationInterfaceComponentArraysExt/DifferentiationInterfaceComponentArraysExt.jl +++ b/ext/DifferentiationInterfaceComponentArraysExt/DifferentiationInterfaceComponentArraysExt.jl @@ -1,26 +1,64 @@ module DifferentiationInterfaceComponentArraysExt using ComponentArrays: ComponentVector -using DifferentiationInterface.DifferentiationTest: Scenario, SCALING_VEC +using DifferentiationInterface.DifferentiationTest: + Scenario, Reference, make_scalar_to_array, make_scalar_to_array!, scalar_to_array_ref +using LinearAlgebra: dot -const SCALING_CVEC = ComponentVector(; a=collect(1:7), b=collect(8:12)) - -function scalar_to_componentvector(x::Number)::ComponentVector - return sin.(SCALING_CVEC .* x) # output size 12 -end +## Vector to scalar function componentvector_to_scalar(x::ComponentVector)::Number return sum(sin, x.a) + sum(cos, x.b) end +componentvector_to_scalar_gradient(x) = ComponentVector(; a=cos.(x.a), b=-sin.(x.b)) + +function componentvector_to_scalar_pushforward(x, dx) + return dot(componentvector_to_scalar_gradient(x), dx) +end + +function componentvector_to_scalar_pullback(x, dy) + return componentvector_to_scalar_gradient(x) .* dy +end + +function componentvector_to_scalar_ref() + return Reference(; + pushforward=componentvector_to_scalar_pushforward, + pullback=componentvector_to_scalar_pullback, + gradient=componentvector_to_scalar_gradient, + ) +end + +## Gather + +const SCALING_CVEC = ComponentVector(; a=collect(1:7), b=collect(8:12)) + function component_scenarios_allocating() return [ - Scenario(scalar_to_componentvector; x=2.0), + Scenario( + make_scalar_to_array(SCALING_CVEC); x=2.0, ref=scalar_to_array_ref(SCALING_CVEC) + ), Scenario( componentvector_to_scalar; x=ComponentVector{Float64}(; a=collect(1:7), b=collect(8:12)), + ref=componentvector_to_scalar_ref(), ), ] end +function component_scenarios_mutating() + return [ + Scenario( + make_scalar_to_array!(SCALING_CVEC); + x=2.0, + y=float.(SCALING_CVEC), + ref=scalar_to_array_ref(SCALING_CVEC), + ), + ] +end + +function component_scenarios() + return vcat(component_scenarios_allocating(), component_scenarios_mutating()) +end + end diff --git a/ext/DifferentiationInterfaceCorrectnessTestExt/DifferentiationInterfaceCorrectnessTestExt.jl b/ext/DifferentiationInterfaceCorrectnessTestExt/DifferentiationInterfaceCorrectnessTestExt.jl deleted file mode 100644 index a92d05c0d..000000000 --- a/ext/DifferentiationInterfaceCorrectnessTestExt/DifferentiationInterfaceCorrectnessTestExt.jl +++ /dev/null @@ -1,337 +0,0 @@ -module DifferentiationInterfaceCorrectnessTestExt - -using ADTypes: AbstractADType -using DifferentiationInterface -using DifferentiationInterface: myisapprox, mysimilar -using DifferentiationInterface.DifferentiationTest: Scenario -import DifferentiationInterface.DifferentiationTest as DT -using ForwardDiff: ForwardDiff -using LinearAlgebra: dot -using Test: @testset, @test -using Zygote: Zygote - -function test_scen_intact(new_scen, scen) - @testset "Scenario intact" begin - @test myisapprox(new_scen.x, scen.x) - @test myisapprox(new_scen.y, scen.y) - @test myisapprox(new_scen.dx, scen.dx) - @test myisapprox(new_scen.dy, scen.dy) - end -end - -## Pushforward - -function test_correctness(ba::AbstractADType, ::typeof(pushforward), scen::Scenario{false}) - (; f, x, y, dx, dy) = new_scen = deepcopy(scen) - dy_true = true_pushforward(f, x, y, dx; mutating=false) - - y_out1, dy_out1 = value_and_pushforward(f, ba, x, dx) - dy_in2 = mysimilar(dy) - y_out2, dy_out2 = value_and_pushforward!!(f, dy_in2, ba, x, dx) - - dy_out3 = pushforward(f, ba, x, dx) - dy_in4 = mysimilar(dy) - dy_out4 = pushforward!!(f, dy_in4, ba, x, dx) - - @testset "Primal value" begin - @test myisapprox(y_out1, y) - @test myisapprox(y_out2, y) - end - @testset "Tangent value" begin - @test myisapprox(dy_out1, dy_true; rtol=1e-3) - @test myisapprox(dy_out2, dy_true; rtol=1e-3) - @test myisapprox(dy_out3, dy_true; rtol=1e-3) - @test myisapprox(dy_out4, dy_true; rtol=1e-3) - end - return test_scen_intact(new_scen, scen) -end - -function test_correctness(ba::AbstractADType, ::typeof(pushforward), scen::Scenario{true}) - (; f, x, y, dx, dy) = new_scen = deepcopy(scen) - f! = f - dy_true = true_pushforward(f!, x, y, dx; mutating=true) - - y_in = mysimilar(y) - dy_in = mysimilar(dy) - y_out, dy_out = value_and_pushforward!!(f!, y_in, dy_in, ba, x, dx) - - @testset "Primal value" begin - @test myisapprox(y_out, y) - end - @testset "Tangent value" begin - @test myisapprox(dy_out, dy_true; rtol=1e-3) - end - return test_scen_intact(new_scen, scen) -end - -## Pullback - -function test_correctness(ba::AbstractADType, ::typeof(pullback), scen::Scenario{false}) - (; f, x, y, dx, dy) = new_scen = deepcopy(scen) - dx_true = true_pullback(f, x, y, dy; mutating=false) - - y_out1, dx_out1 = value_and_pullback(f, ba, x, dy) - dx_in2 = mysimilar(dx) - y_out2, dx_out2 = value_and_pullback!!(f, dx_in2, ba, x, dy) - - dx_out3 = pullback(f, ba, x, dy) - dx_in4 = mysimilar(dx) - dx_out4 = pullback!!(f, dx_in4, ba, x, dy) - - @testset "Primal value" begin - @test myisapprox(y_out1, y) - @test myisapprox(y_out2, y) - end - @testset "Cotangent value" begin - @test myisapprox(dx_out1, dx_true; rtol=1e-3) - @test myisapprox(dx_out2, dx_true; rtol=1e-3) - @test myisapprox(dx_out3, dx_true; rtol=1e-3) - @test myisapprox(dx_out4, dx_true; rtol=1e-3) - end - return test_scen_intact(new_scen, scen) -end - -function test_correctness(ba::AbstractADType, ::typeof(pullback), scen::Scenario{true}) - (; f, x, y, dx, dy) = new_scen = deepcopy(scen) - f! = f - dx_true = true_pullback(f, x, y, dy; mutating=true) - - y_in = mysimilar(y) - dx_in = mysimilar(dx) - y_out, dx_out = value_and_pullback!!(f!, y_in, dx_in, ba, x, dy) - - @testset "Primal value" begin - @test myisapprox(y_out, y) - end - @testset "Cotangent value" begin - @test myisapprox(dx_out, dx_true; rtol=1e-3) - end - return test_scen_intact(new_scen, scen) -end - -## Derivative - -function test_correctness(ba::AbstractADType, ::typeof(derivative), scen::Scenario{false}) - (; f, x, y, dx, dy) = new_scen = deepcopy(scen) - der_true = ForwardDiff.derivative(f, x) - - y_out1, der_out1 = value_and_derivative(f, ba, x) - der_in2 = mysimilar(dy) - y_out2, der_out2 = value_and_derivative!!(f, der_in2, ba, x) - - der_out3 = derivative(f, ba, x) - der_in4 = mysimilar(dy) - der_out4 = derivative!!(f, der_in4, ba, x) - - @testset "Primal value" begin - @test myisapprox(y_out1, y) - @test myisapprox(y_out2, y) - end - @testset "Derivative value" begin - @test myisapprox(der_out1, der_true; rtol=1e-3) - @test myisapprox(der_out2, der_true; rtol=1e-3) - @test myisapprox(der_out3, der_true; rtol=1e-3) - @test myisapprox(der_out4, der_true; rtol=1e-3) - end - return test_scen_intact(new_scen, scen) -end - -function test_correctness(ba::AbstractADType, ::typeof(derivative), scen::Scenario{true}) - (; f, x, y, dx, dy) = new_scen = deepcopy(scen) - f! = f - der_true = ForwardDiff.derivative(f!, y, x) - - y_in = mysimilar(y) - der_in = mysimilar(dy) - y_out, der_out = value_and_derivative!!(f!, y_in, der_in, ba, x) - - @testset "Primal value" begin - @test myisapprox(y_out, y) - end - @testset "Derivative value" begin - @test myisapprox(der_out, der_true; rtol=1e-3) - end - return test_scen_intact(new_scen, scen) -end - -## Gradient - -function test_correctness(ba::AbstractADType, ::typeof(gradient), scen::Scenario{false}) - (; f, x, y, dx, dy) = new_scen = deepcopy(scen) - grad_true = if x isa Number - ForwardDiff.derivative(f, x) - else - try - ForwardDiff.gradient(f, x) - catch e - only(Zygote.gradient(f, x)) - end - end - - y_out1, grad_out1 = value_and_gradient(f, ba, x) - grad_in2 = mysimilar(dx) - y_out2, grad_out2 = value_and_gradient!!(f, grad_in2, ba, x) - - grad_out3 = gradient(f, ba, x) - grad_in4 = mysimilar(dx) - grad_out4 = gradient!!(f, grad_in4, ba, x) - - @testset "Primal value" begin - @test myisapprox(y_out1, y) - @test myisapprox(y_out2, y) - end - @testset "Gradient value" begin - @test myisapprox(grad_out1, grad_true; rtol=1e-3) - @test myisapprox(grad_out2, grad_true; rtol=1e-3) - @test myisapprox(grad_out3, grad_true; rtol=1e-3) - @test myisapprox(grad_out4, grad_true; rtol=1e-3) - end - return test_scen_intact(new_scen, scen) -end - -## Jacobian - -function test_correctness(ba::AbstractADType, ::typeof(jacobian), scen::Scenario{false}) - (; f, x, y) = new_scen = deepcopy(scen) - jac_true = ForwardDiff.jacobian(f, x) - - y_out1, jac_out1 = value_and_jacobian(f, ba, x) - jac_in2 = mysimilar(jac_true) - y_out2, jac_out2 = value_and_jacobian!!(f, jac_in2, ba, x) - - jac_out3 = jacobian(f, ba, x) - jac_in4 = mysimilar(jac_true) - jac_out4 = jacobian!!(f, jac_in4, ba, x) - - @testset "Primal value" begin - @test myisapprox(y_out1, y) - @test myisapprox(y_out2, y) - end - @testset "Jacobian value" begin - @test myisapprox(jac_out1, jac_true; rtol=1e-3) - @test myisapprox(jac_out2, jac_true; rtol=1e-3) - @test myisapprox(jac_out3, jac_true; rtol=1e-3) - @test myisapprox(jac_out4, jac_true; rtol=1e-3) - end - return test_scen_intact(new_scen, scen) -end - -function test_correctness(ba::AbstractADType, ::typeof(jacobian), scen::Scenario{true}) - (; f, x, y) = new_scen = deepcopy(scen) - f! = f - jac_true = ForwardDiff.jacobian(f!, y, x) - - y_in = mysimilar(y) - jac_in = mysimilar(jac_true) - y_out, jac_out = value_and_jacobian!!(f!, y_in, jac_in, ba, x) - - @testset "Primal value" begin - @test myisapprox(y_out, y) - end - @testset "Jacobian value" begin - @test myisapprox(jac_out, jac_true; rtol=1e-3) - end - return test_scen_intact(new_scen, scen) -end - -## Second derivative - -function test_correctness(ba::AbstractADType, ::typeof(second_derivative), scen::Scenario) - (; f, x) = deepcopy(scen) - der2_true = ForwardDiff.derivative(z -> ForwardDiff.derivative(f, z), x) - - der2_out1 = second_derivative(f, ba, x) - - @testset "Second derivative value" begin - @test myisapprox(der2_out1, der2_true; rtol=1e-3) - end -end - -## Hessian-vector product - -function test_correctness(ba::AbstractADType, ::typeof(hvp), scen::Scenario) - (; f, x, dx) = deepcopy(scen) - hess_true = if x isa Number - ForwardDiff.derivative(z -> ForwardDiff.derivative(f, z), x) - else - ForwardDiff.hessian(f, x) - end - hvp_true = if x isa Number - hess_true * dx - else - reshape((hess_true * vec(dx)), size(x)) - end - - hvp_out1 = hvp(f, ba, x, dx) - - @testset "Hessian-vector product value" begin - @test myisapprox(hvp_out1, hvp_true; rtol=1e-3) - end -end - -## Hessian - -function test_correctness(ba::AbstractADType, ::typeof(hessian), scen::Scenario) - (; f, x, y) = deepcopy(scen) - hess_true = ForwardDiff.hessian(f, x) - - hess_out1 = hessian(f, ba, x) - - @testset "Hessian value" begin - @test myisapprox(hess_out1, hess_true; rtol=1e-3) - end -end - -## Utils - -function true_pushforward(f, x::Number, y::Number, dx; mutating) - return ForwardDiff.derivative(f, x) * dx -end - -function true_pushforward(f, x::Number, y::AbstractArray, dx; mutating) - if mutating - return ForwardDiff.derivative(f, deepcopy(y), x) .* dx - else - return ForwardDiff.derivative(f, x) .* dx - end -end - -function true_pushforward(f, x::AbstractArray, y::Number, dx; mutating) - return dot(Zygote.gradient(f, x)[1], dx) -end - -function true_pushforward(f, x::AbstractArray, y::AbstractArray, dx; mutating) - if mutating - return reshape(ForwardDiff.jacobian(f, deepcopy(y), x) * vec(dx), size(y)) - else - return reshape(ForwardDiff.jacobian(f, x) * vec(dx), size(y)) - end -end - -function true_pullback(f, x::Number, y::Number, dy; mutating) - return ForwardDiff.derivative(f, x) * dy -end - -function true_pullback(f, x::Number, y::AbstractArray, dy; mutating) - if mutating - return dot(ForwardDiff.derivative(f, deepcopy(y), x), dy) - else - return dot(ForwardDiff.derivative(f, x), dy) - end -end - -function true_pullback(f, x::AbstractArray, y::Number, dy; mutating) - return Zygote.gradient(f, x)[1] .* dy -end - -function true_pullback(f, x::AbstractArray, y::AbstractArray, dy; mutating) - if mutating - return reshape( - transpose(ForwardDiff.jacobian(f, deepcopy(y), x)) * vec(dy), size(x) - ) - else - return reshape(transpose(ForwardDiff.jacobian(f, x)) * vec(dy), size(x)) - end -end - -end # module diff --git a/ext/DifferentiationInterfaceFastDifferentiationExt/allocating.jl b/ext/DifferentiationInterfaceFastDifferentiationExt/allocating.jl index ec0ed3ccf..31728a10d 100644 --- a/ext/DifferentiationInterfaceFastDifferentiationExt/allocating.jl +++ b/ext/DifferentiationInterfaceFastDifferentiationExt/allocating.jl @@ -8,8 +8,8 @@ function DI.prepare_pushforward(f, ::AutoFastDifferentiation, x) end y_var = f(x_var) - x_vec_var = x_var isa Number ? [x_var] : x_var - y_vec_var = y_var isa Number ? [y_var] : y_var + x_vec_var = x_var isa Number ? [x_var] : vec(x_var) + y_vec_var = y_var isa Number ? [y_var] : vec(y_var) jv_vec_var, v_vec_var = jacobian_times_v(y_vec_var, x_vec_var) jvp_exe = make_function(jv_vec_var, [x_vec_var; v_vec_var]; in_place=false) return jvp_exe @@ -24,7 +24,7 @@ function DI.value_and_pushforward( if y isa Number return y, only(jv_vec) else - return y, jv_vec + return y, reshape(jv_vec, size(y)) end end diff --git a/ext/DifferentiationInterfaceJETExt/DifferentiationInterfaceJETExt.jl b/ext/DifferentiationInterfaceJETExt/DifferentiationInterfaceJETExt.jl index 6cc669b7c..6f98577ef 100644 --- a/ext/DifferentiationInterfaceJETExt/DifferentiationInterfaceJETExt.jl +++ b/ext/DifferentiationInterfaceJETExt/DifferentiationInterfaceJETExt.jl @@ -94,7 +94,7 @@ end function test_jet(ba::AbstractADType, ::typeof(jacobian), scen::Scenario{false};) (; f, x, y) = deepcopy(scen) extras = prepare_jacobian(f, ba, x) - jac_in = similar(x, length(y), length(x)) + jac_in = Matrix{eltype(y)}(undef, length(y), length(x)) @test_opt value_and_jacobian!!(f, jac_in, ba, x, extras) @test_opt value_and_jacobian(f, ba, x, extras) @@ -106,7 +106,7 @@ function test_jet(ba::AbstractADType, ::typeof(jacobian), scen::Scenario{true};) f! = f extras = prepare_jacobian(f!, ba, y, x) y_in = mysimilar(y) - jac_in = similar(x, length(y), length(x)) + jac_in = Matrix{eltype(y)}(undef, length(y), length(x)) @test_opt value_and_jacobian!!(f!, y_in, jac_in, ba, x, extras) return nothing diff --git a/ext/DifferentiationInterfaceJLArraysExt/DifferentiationInterfaceJLArraysExt.jl b/ext/DifferentiationInterfaceJLArraysExt/DifferentiationInterfaceJLArraysExt.jl index 08361e006..95df6af0d 100644 --- a/ext/DifferentiationInterfaceJLArraysExt/DifferentiationInterfaceJLArraysExt.jl +++ b/ext/DifferentiationInterfaceJLArraysExt/DifferentiationInterfaceJLArraysExt.jl @@ -2,37 +2,51 @@ module DifferentiationInterfaceJLArraysExt using DifferentiationInterface.DifferentiationTest: Scenario, - vector_to_scalar, - matrix_to_scalar, + Reference, + make_scalar_to_array, + scalar_to_array_ref, + array_to_scalar, + array_to_scalar_ref, vector_to_vector, + vector_to_vector_ref, vector_to_matrix, + vector_to_matrix_ref, matrix_to_vector, - matrix_to_matrix + matrix_to_vector_ref, + matrix_to_matrix, + matrix_to_matrix_ref using JLArrays const SCALING_JLVEC = jl(Vector(1:12)) const SCALING_JLMAT = jl(Matrix((1:3) .* transpose(1:4))) -function scalar_to_jlvector(x::Number)::JLArray{<:Any,1} - return sin.(SCALING_JLVEC .* x) # output size 12 -end - -function scalar_to_jlmatrix(x::Number)::JLArray{<:Any,2} - return sin.(SCALING_JLMAT .* x) # output size (3, 4) -end - function gpu_scenarios_allocating() - scenarios = [ - Scenario(scalar_to_jlvector; x=2.0), - Scenario(scalar_to_jlmatrix; x=2.0), - Scenario(vector_to_scalar; x=jl(Vector{Float64}(1:12))), - Scenario(matrix_to_scalar; x=jl(Matrix{Float64}(reshape(1:12, 3, 4)))), - Scenario(vector_to_vector; x=jl(Vector{Float64}(1:12))), - Scenario(vector_to_matrix; x=jl(Vector{Float64}(1:12))), - Scenario(matrix_to_vector; x=jl(Matrix{Float64}(reshape(1:12, 3, 4)))), - Scenario(matrix_to_matrix; x=jl(Matrix{Float64}(reshape(1:12, 3, 4)))), + return [ + Scenario( + make_scalar_to_array(SCALING_JLVEC); + x=2.0, + ref=scalar_to_array_ref(SCALING_JLVEC), + ), + Scenario( + make_scalar_to_array(SCALING_JLMAT); + x=2.0, + ref=scalar_to_array_ref(SCALING_JLMAT), + ), + Scenario(array_to_scalar; x=jl(float.(1:12)), ref=array_to_scalar_ref()), + Scenario( + array_to_scalar; x=jl(float.(reshape(1:12, 3, 4))), ref=array_to_scalar_ref() + ), + Scenario(vector_to_vector; x=jl(float.(1:12)), ref=vector_to_vector_ref()), + Scenario(vector_to_matrix; x=jl(float.(1:12)), ref=vector_to_matrix_ref()), + Scenario( + matrix_to_vector; x=jl(float.(reshape(1:12, 3, 4))), ref=matrix_to_vector_ref() + ), + Scenario( + matrix_to_matrix; x=jl(float.(reshape(1:12, 3, 4))), ref=matrix_to_matrix_ref() + ), ] - return scenarios end +gpu_scenarios() = gpu_scenarios_allocating() + end diff --git a/ext/DifferentiationInterfacePolyesterForwardDiffExt/allocating.jl b/ext/DifferentiationInterfacePolyesterForwardDiffExt/allocating.jl index 913f96f51..60a3b3551 100644 --- a/ext/DifferentiationInterfacePolyesterForwardDiffExt/allocating.jl +++ b/ext/DifferentiationInterfacePolyesterForwardDiffExt/allocating.jl @@ -18,27 +18,35 @@ end ## Gradient function DI.value_and_gradient!!( - f, grad::AbstractArray, ::AutoPolyesterForwardDiff{C}, x::AbstractArray, extras::Nothing + f, + grad::AbstractVector, + ::AutoPolyesterForwardDiff{C}, + x::AbstractVector, + extras::Nothing, ) where {C} threaded_gradient!(f, grad, x, Chunk{C}()) return f(x), grad end function DI.gradient!!( - f, grad::AbstractArray, ::AutoPolyesterForwardDiff{C}, x::AbstractArray, extras::Nothing + f, + grad::AbstractVector, + ::AutoPolyesterForwardDiff{C}, + x::AbstractVector, + extras::Nothing, ) where {C} threaded_gradient!(f, grad, x, Chunk{C}()) return grad end function DI.value_and_gradient( - f, backend::AutoPolyesterForwardDiff, x::AbstractArray, extras::Nothing + f, backend::AutoPolyesterForwardDiff, x::AbstractVector, extras::Nothing ) return DI.value_and_gradient!!(f, mysimilar(x), backend, x, extras) end function DI.gradient( - f, backend::AutoPolyesterForwardDiff, x::AbstractArray, extras::Nothing + f, backend::AutoPolyesterForwardDiff, x::AbstractVector, extras::Nothing ) return DI.gradient!!(f, mysimilar(x), backend, x, extras) end diff --git a/ext/DifferentiationInterfaceStaticArraysExt/DifferentiationInterfaceStaticArraysExt.jl b/ext/DifferentiationInterfaceStaticArraysExt/DifferentiationInterfaceStaticArraysExt.jl index 7dafa6d85..a59691921 100644 --- a/ext/DifferentiationInterfaceStaticArraysExt/DifferentiationInterfaceStaticArraysExt.jl +++ b/ext/DifferentiationInterfaceStaticArraysExt/DifferentiationInterfaceStaticArraysExt.jl @@ -2,37 +2,53 @@ module DifferentiationInterfaceStaticArraysExt using DifferentiationInterface.DifferentiationTest: Scenario, - vector_to_scalar, - matrix_to_scalar, + Reference, + make_scalar_to_array, + scalar_to_array_ref, + array_to_scalar, + array_to_scalar_ref, vector_to_vector, + vector_to_vector_ref, vector_to_matrix, + vector_to_matrix_ref, matrix_to_vector, - matrix_to_matrix + matrix_to_vector_ref, + matrix_to_matrix, + matrix_to_matrix_ref using StaticArrays const SCALING_SVEC = SVector{12}(1:12) const SCALING_SMAT = SMatrix{3,4}((1:3) .* transpose(1:4)) -function scalar_to_svector(x::Number)::SVector - return sin.(SCALING_SVEC .* x) # output size 12 -end - -function scalar_to_smatrix(x::Number)::SMatrix - return sin.(SCALING_SMAT .* x) # output size (3, 4) -end - function static_scenarios_allocating() - scenarios = [ - Scenario(scalar_to_svector; x=2.0), - Scenario(scalar_to_smatrix; x=2.0), - Scenario(vector_to_scalar; x=SVector{12,Float64}(1:12)), - Scenario(matrix_to_scalar; x=SMatrix{3,4,Float64}(reshape(1:12, 3, 4))), - Scenario(vector_to_vector; x=SVector{12,Float64}(1:12)), - Scenario(vector_to_matrix; x=SVector{12,Float64}(1:12)), - Scenario(matrix_to_vector; x=SMatrix{3,4,Float64}(reshape(1:12, 3, 4))), - Scenario(matrix_to_matrix; x=SMatrix{3,4,Float64}(reshape(1:12, 3, 4))), + return [ + Scenario( + make_scalar_to_array(SCALING_SVEC); x=2.0, ref=scalar_to_array_ref(SCALING_SVEC) + ), + Scenario( + make_scalar_to_array(SCALING_SMAT); x=2.0, ref=scalar_to_array_ref(SCALING_SMAT) + ), + Scenario(array_to_scalar; x=SVector{12,Float64}(1:12), ref=array_to_scalar_ref()), + Scenario( + array_to_scalar; + x=SMatrix{3,4,Float64}(reshape(1:12, 3, 4)), + ref=array_to_scalar_ref(), + ), + Scenario(vector_to_vector; x=SVector{12,Float64}(1:12), ref=vector_to_vector_ref()), + Scenario(vector_to_matrix; x=SVector{12,Float64}(1:12), ref=vector_to_matrix_ref()), + Scenario( + matrix_to_vector; + x=SMatrix{3,4,Float64}(reshape(1:12, 3, 4)), + ref=matrix_to_vector_ref(), + ), + Scenario( + matrix_to_matrix; + x=SMatrix{3,4,Float64}(reshape(1:12, 3, 4)), + ref=matrix_to_matrix_ref(), + ), ] - return scenarios end +static_scenarios() = static_scenarios_allocating() + end diff --git a/src/DifferentiationTest/DifferentiationTest.jl b/src/DifferentiationTest/DifferentiationTest.jl index 6e8bb39a3..614bb88f1 100644 --- a/src/DifferentiationTest/DifferentiationTest.jl +++ b/src/DifferentiationTest/DifferentiationTest.jl @@ -19,6 +19,7 @@ using ..DifferentiationInterface: AutoTaped, inner, mode, + myisapprox, mysimilar, mysimilar_random, myzero, @@ -29,22 +30,27 @@ using ..DifferentiationInterface: supports_pullback using DocStringExtensions using Functors: @functor, fleaves, fmap -using LinearAlgebra: dot +using LinearAlgebra: Diagonal, dot using Test: @testset, @test -include("scenario.jl") -include("compatibility.jl") -include("scenarios_default.jl") -include("scenarios_weird_arrays.jl") -include("scenarios_nested.jl") -include("zero.jl") -include("printing.jl") -include("benchmark.jl") -include("test_call_count.jl") -include("test_error_free.jl") -include("test_differentiation.jl") +include("scenarios/scenario.jl") +include("scenarios/default.jl") +include("scenarios/weird_arrays.jl") +include("scenarios/nested.jl") +# include("scenarios/layer.jl") -export Scenario, default_scenarios, weird_array_scenarios, nested_scenarios +include("utils/zero.jl") +include("utils/compatibility.jl") +include("utils/printing.jl") + +include("tests/correctness.jl") +include("tests/call_count.jl") +include("tests/benchmark.jl") +include("tests/test.jl") + +export Scenario +export default_scenarios +export weird_array_scenarios, nested_scenarios export BenchmarkData, record! export all_operators, test_differentiation diff --git a/src/DifferentiationTest/scenario.jl b/src/DifferentiationTest/scenario.jl deleted file mode 100644 index 2558ef3f3..000000000 --- a/src/DifferentiationTest/scenario.jl +++ /dev/null @@ -1,56 +0,0 @@ - -""" - Scenario{mutating} - -Store a testing scenario composed of a function and its input + output + tangents. - -# Fields - -$(TYPEDFIELDS) -""" -struct Scenario{mutating,F,X,Y,DX,DY} - "function" - f::F - "input" - x::X - "output" - y::Y - "pushforward seed" - dx::DX - "pullback seed" - dy::DY -end - -function Scenario{mutating}(f::F, x::X, y::Y, dx::DX, dy::DY) where {mutating,F,X,Y,DX,DY} - return Scenario{mutating,F,X,Y,DX,DY}(f, x, y, dx, dy) -end - -function Base.string(scen::Scenario{mutating}) where {mutating} - return "$(string(scen.f)): $(typeof(scen.x)) -> $(typeof(scen.y))" -end - -is_mutating(::Scenario{mutating}) where {mutating} = mutating - -## Scenario constructors - -function Scenario(f, x) - y = f(x) - dx = mysimilar_random(x) - dy = mysimilar_random(y) - return Scenario{false}(f, x, y, dx, dy) -end - -function Scenario(f!, y, x) - f!(y, x) - dx = mysimilar_random(x) - dy = mysimilar_random(y) - return Scenario{true}(f!, x, y, dx, dy) -end - -function Scenario(f; x, y=nothing) - if isnothing(y) - return Scenario(f, x) - else - return Scenario(f, y, x) - end -end diff --git a/src/DifferentiationTest/scenarios/default.jl b/src/DifferentiationTest/scenarios/default.jl new file mode 100644 index 000000000..7626d4ee8 --- /dev/null +++ b/src/DifferentiationTest/scenarios/default.jl @@ -0,0 +1,254 @@ +#= +Constraints on the scenarios: +- non-allocating whenever possible +- type-stable +- GPU-compatible (no scalar indexing) +- vary shapes to be tricky +=# + +first_half(v::AbstractVector) = @view v[1:(length(v) ÷ 2)] +second_half(v::AbstractVector) = @view v[(length(v) ÷ 2 + 1):end] + +top_half(M::AbstractMatrix) = @view M[1:(size(M, 2) ÷ 2), :] +bottom_half(M::AbstractMatrix) = @view M[(size(M, 2) ÷ 2 + 1):end, :] + +left_half(M::AbstractMatrix) = @view M[:, 1:(size(M, 2) ÷ 2)] +right_half(M::AbstractMatrix) = @view M[:, (size(M, 2) ÷ 2 + 1):end] + +## Scalar to scalar + +scalar_to_scalar(x::Number)::Number = sin(x) + +scalar_to_scalar_derivative(x) = cos(x) +scalar_to_scalar_second_derivative(x) = -sin(x) +scalar_to_scalar_pushforward(x, dx) = scalar_to_scalar_derivative(x) * dx +scalar_to_scalar_pullback(x, dy) = scalar_to_scalar_derivative(x) * dy +scalar_to_scalar_gradient(x) = scalar_to_scalar_derivative(x) +scalar_to_scalar_hvp(x, v) = scalar_to_scalar_second_derivative(x) * v + +function scalar_to_scalar_ref() + return Reference(; + pushforward=scalar_to_scalar_pushforward, + pullback=scalar_to_scalar_pullback, + derivative=scalar_to_scalar_derivative, + gradient=scalar_to_scalar_gradient, + second_derivative=scalar_to_scalar_second_derivative, + hvp=scalar_to_scalar_hvp, + ) +end + +## Scalar to array + +_scalar_to_array(x::Number, scaling::AbstractArray)::AbstractArray = sin.(scaling .* x) + +function _scalar_to_array!(y::AbstractArray, x::Number, scaling::AbstractArray)::Nothing + y .= sin.(scaling .* x) + return nothing +end + +function make_scalar_to_array(scaling::AbstractArray) + scalar_to_array(x::Number) = _scalar_to_array(x, scaling) + return scalar_to_array +end + +function make_scalar_to_array!(scaling::AbstractArray) + scalar_to_array!(y::AbstractArray, x::Number) = _scalar_to_array!(y, x, scaling) + return scalar_to_array! +end + +scalar_to_array_derivative(x, scaling) = scaling .* cos.(scaling .* x) +scalar_to_array_second_derivative(x, scaling) = -(scaling .^ 2) .* sin.(scaling .* x) +scalar_to_array_pushforward(x, dx, scaling) = scalar_to_array_derivative(x, scaling) .* dx +scalar_to_array_pullback(x, dy, scaling) = dot(scalar_to_array_derivative(x, scaling), dy) + +function scalar_to_array_ref(scaling) + return Reference(; + pushforward=(args...) -> scalar_to_array_pushforward(args..., scaling), + pullback=(args...) -> scalar_to_array_pullback(args..., scaling), + derivative=(args...) -> scalar_to_array_derivative(args..., scaling), + second_derivative=(args...) -> scalar_to_array_second_derivative(args..., scaling), + ) +end + +## Array to scalar + +array_to_scalar(x::AbstractArray)::Number = sum(sin, x) + +array_to_scalar_gradient(x) = cos.(x) +array_to_scalar_hvp(x, v) = -sin.(x) .* v +array_to_scalar_pushforward(x, dx) = dot(array_to_scalar_gradient(x), dx) +array_to_scalar_pullback(x, dy) = array_to_scalar_gradient(x) .* dy +array_to_scalar_hessian(x) = Diagonal(-sin.(vec(x))) + +function array_to_scalar_ref() + return Reference(; + pushforward=array_to_scalar_pushforward, + pullback=array_to_scalar_pullback, + gradient=array_to_scalar_gradient, + hvp=array_to_scalar_hvp, + hessian=array_to_scalar_hessian, + ) +end + +## Array to array + +vector_to_vector(x::AbstractVector)::AbstractVector = vcat(sin.(x), cos.(x)) + +function vector_to_vector!(y::AbstractVector, x::AbstractVector) + y[1:length(x)] .= sin.(x) + y[(length(x) + 1):(2length(x))] .= cos.(x) + return nothing +end + +vector_to_vector_pushforward(x, dx) = vcat(cos.(x) .* dx, -sin.(x) .* dx) +vector_to_vector_pullback(x, dy) = cos.(x) .* first_half(dy) .- sin.(x) .* second_half(dy) +vector_to_vector_jacobian(x) = vcat(Diagonal(cos.(x)), Diagonal(-sin.(x))) + +function vector_to_vector_ref() + return Reference(; + pushforward=vector_to_vector_pushforward, + pullback=vector_to_vector_pullback, + jacobian=vector_to_vector_jacobian, + ) +end + +vector_to_matrix(x::AbstractVector)::AbstractMatrix = hcat(sin.(x), cos.(x)) + +function vector_to_matrix!(y::AbstractMatrix, x::AbstractVector) + y[:, 1] .= sin.(x) + y[:, 2] .= cos.(x) + return nothing +end + +vector_to_matrix_pushforward(x, dx) = hcat(cos.(x) .* dx, -sin.(x) .* dx) +vector_to_matrix_pullback(x, dy) = cos.(x) .* dy[:, 1] .- sin.(x) .* dy[:, 2] +vector_to_matrix_jacobian(x) = vcat(Diagonal(cos.(x)), Diagonal(-sin.(x))) + +function vector_to_matrix_ref() + return Reference(; + pushforward=vector_to_matrix_pushforward, + pullback=vector_to_matrix_pullback, + jacobian=vector_to_matrix_jacobian, + ) +end + +matrix_to_vector(x::AbstractMatrix)::AbstractVector = vcat(vec(sin.(x)), vec(cos.(x))) + +function matrix_to_vector!(y::AbstractVector, x::AbstractMatrix) + n = length(x) + y[1:n] .= sin.(getindex.(Ref(x), 1:n)) + y[(n + 1):(2n)] .= cos.(getindex.(Ref(x), 1:n)) + return nothing +end + +matrix_to_vector_pushforward(x, dx) = vcat(vec(cos.(x) .* dx), vec(-sin.(x) .* dx)) +function matrix_to_vector_pullback(x, dy) + return cos.(x) .* reshape(first_half(dy), size(x)) .- + sin.(x) .* reshape(second_half(dy), size(x)) +end +matrix_to_vector_jacobian(x) = vcat(Diagonal(vec(cos.(x))), Diagonal(vec(-sin.(x)))) + +function matrix_to_vector_ref() + return Reference(; + pushforward=matrix_to_vector_pushforward, + pullback=matrix_to_vector_pullback, + jacobian=matrix_to_vector_jacobian, + ) +end + +matrix_to_matrix(x::AbstractMatrix)::AbstractMatrix = hcat(vec(sin.(x)), vec(cos.(x))) + +function matrix_to_matrix!(y::AbstractMatrix, x::AbstractMatrix) + n = length(x) + y[:, 1] .= sin.(getindex.(Ref(x), 1:n)) + y[:, 2] .= cos.(getindex.(Ref(x), 1:n)) + return nothing +end + +matrix_to_matrix_pushforward(x, dx) = hcat(vec(cos.(x) .* dx), vec(-sin.(x) .* dx)) +function matrix_to_matrix_pullback(x, dy) + return cos.(x) .* reshape(dy[:, 1], size(x)) .- sin.(x) .* reshape(dy[:, 2], size(x)) +end +matrix_to_matrix_jacobian(x) = vcat(Diagonal(vec(cos.(x))), Diagonal(vec(-sin.(x)))) + +function matrix_to_matrix_ref() + return Reference(; + pushforward=matrix_to_matrix_pushforward, + pullback=matrix_to_matrix_pullback, + jacobian=matrix_to_matrix_jacobian, + ) +end + +## Gather + +const SCALING_VEC = Vector(1:12) +const SCALING_MAT = Matrix((1:3) .* transpose(1:4)) + +function default_scenarios_allocating() + return [ + Scenario(scalar_to_scalar; x=2.0, ref=scalar_to_scalar_ref()), + Scenario( + make_scalar_to_array(SCALING_VEC); x=2.0, ref=scalar_to_array_ref(SCALING_VEC) + ), + Scenario( + make_scalar_to_array(SCALING_MAT); x=2.0, ref=scalar_to_array_ref(SCALING_MAT) + ), + Scenario(array_to_scalar; x=float.(1:12), ref=array_to_scalar_ref()), + Scenario(array_to_scalar; x=float.(reshape(1:12, 3, 4)), ref=array_to_scalar_ref()), + Scenario(vector_to_vector; x=float.(1:12), ref=vector_to_vector_ref()), + Scenario(vector_to_matrix; x=float.(1:12), ref=vector_to_matrix_ref()), + Scenario( + matrix_to_vector; x=float.(reshape(1:12, 3, 4)), ref=matrix_to_vector_ref() + ), + Scenario( + matrix_to_matrix; x=float.(reshape(1:12, 3, 4)), ref=matrix_to_matrix_ref() + ), + ] +end + +function default_scenarios_mutating() + return [ + Scenario( + make_scalar_to_array!(SCALING_VEC); + x=2.0, + y=zeros(12), + ref=scalar_to_array_ref(SCALING_VEC), + ), + Scenario( + make_scalar_to_array!(SCALING_MAT); + x=2.0, + y=zeros(3, 4), + ref=scalar_to_array_ref(SCALING_MAT), + ), + Scenario( + vector_to_vector!; x=float.(1:12), y=zeros(24), ref=vector_to_vector_ref() + ), + Scenario( + vector_to_matrix!; x=float.(1:12), y=zeros(12, 2), ref=vector_to_matrix_ref() + ), + Scenario( + matrix_to_vector!; + x=float.(reshape(1:12, 3, 4)), + y=zeros(24), + ref=matrix_to_vector_ref(), + ), + Scenario( + matrix_to_matrix!; + x=float.(reshape(1:12, 3, 4)), + y=zeros(12, 2), + ref=matrix_to_matrix_ref(), + ), + ] +end + +""" + default_scenarios() + +Create a vector of [`Scenario`](@ref)s for testing differentiation. +""" +function default_scenarios() + return vcat( + default_scenarios_allocating(), # + default_scenarios_mutating(), # + ) +end diff --git a/src/DifferentiationTest/scenarios/layer.jl b/src/DifferentiationTest/scenarios/layer.jl new file mode 100644 index 000000000..efea07b54 --- /dev/null +++ b/src/DifferentiationTest/scenarios/layer.jl @@ -0,0 +1,18 @@ +@kwdef struct Layer{W,B,A} + w::W + b::B + σ::A = nothing +end + +@functor Layer + +(l::Layer{<:Number,<:Number,<:Nothing})(x::Number) = l.w * x + l.b +(l::Layer{<:Number,<:Number})(x::Number) = l.σ(l.w * x + l.b) + +(l::Layer{<:AbstractMatrix,<:AbstractVector,<:Nothing})(x::AbstractVector) = l.w * x + l.b +(l::Layer{<:AbstractMatrix,<:AbstractVector})(x::AbstractVector) = l.σ.(l.w * x + l.b) + +call_layer(l::Layer{<:Number,<:Number}) = l(3.0) +call_layer(l::Layer{<:AbstractMatrix,<:AbstractVector}) = l(3 * ones(size(l.w, 2))) + +sum_call_layer(l) = sum(call_layer(l)) diff --git a/src/DifferentiationTest/scenarios/nested.jl b/src/DifferentiationTest/scenarios/nested.jl new file mode 100644 index 000000000..3d940a290 --- /dev/null +++ b/src/DifferentiationTest/scenarios/nested.jl @@ -0,0 +1,80 @@ +recursive_norm(x::Number) = abs2(x) +recursive_norm(x::AbstractArray) = sum(abs2, x) +recursive_norm(x) = sum(recursive_norm, fleaves(x)) + +recursive_norm_gradient(x) = fmap(_x -> 2_x, x) +recursive_norm_hvp(x) = fmap(_x -> 2 * one(x), x) +recursive_norm_pushforward(x, dx) = fmap(dot, recursive_norm_gradient(x), dx) +recursive_norm_pullback(x, dy) = fmap(_x -> dy * _x, recursive_norm_gradient(x)) + +function recursive_norm_ref() + return Reference(; + pushforward=recursive_norm_pushforward, + pullback=recursive_norm_pullback, + gradient=recursive_norm_gradient, + hvp=recursive_norm_hvp, + ) +end + +function nested(x::Number) + return (a=[x^2, x^3], b=([sin(x);;],), c=1.0) +end + +nested_derivative(x) = (a=[2x, 3x^2], b=([cos(x);;],), c=0.0) +nested_pushforward(x, dx) = (a=[2x * dx, 3x^2 * dx], b=([cos(x) * dx;;],), c=0.0) +nested_pullback(x, dy) = (a=[2x * dy, 3x^2 * dy], b=([cos(x) * dy;;],), c=0.0) +nested_second_derivative(x) = (a=[2, 6x], b=([-sin(x);;],), c=0.0) + +function nested_ref() + return Reference(; + pushforward=nested_pushforward, + pullback=nested_pullback, + derivative=nested_derivative, + second_derivative=nested_second_derivative, + ) +end + +function nested_immutables(x::Number) + return merge(nested(x), (; d=cos(x))) +end + +function nested_immutables_derivative(x) + return merge(nested_derivative(x), (; d=-sin(x))) +end + +function nested_immutables_pushforward(x, dx) + return merge(nested_pushforward(x, dx), (; d=-sin(x) * dx)) +end + +function nested_immutables_pullback(x, dy) + return merge(nested_pullback(x, dy), (; d=-sin(x) * dy)) +end + +function nested_immutables_second_derivative(x) + return merge(nested_second_derivative(x), (; d=-cos(x))) +end + +function nested_immutables_ref() + return Reference(; + derivative=nested_immutables_derivative, + pushforward=nested_immutables_pushforward, + pullback=nested_immutables_pullback, + second_derivative=nested_immutables_second_derivative, + ) +end + +function nested_scenarios(; immutables=true) + scenarios = [ + Scenario(nested; x=2.0, ref=nested_ref()), + Scenario(recursive_norm; x=nested(2.0), ref=recursive_norm_ref()), + ] + scenarios_immutables = [ + Scenario(nested_immutables; x=2.0, ref=nested_immutables_ref()), + Scenario(recursive_norm; x=nested_immutables(2.0), ref=recursive_norm_ref()), + ] + if immutables + return vcat(scenarios, scenarios_immutables) + else + return scenarios + end +end diff --git a/src/DifferentiationTest/scenarios/scenario.jl b/src/DifferentiationTest/scenarios/scenario.jl new file mode 100644 index 000000000..e1be1e9aa --- /dev/null +++ b/src/DifferentiationTest/scenarios/scenario.jl @@ -0,0 +1,84 @@ +""" + Reference + +Store the ground truth operators for a [`Scenario`](@ref). + +# Fields + +$(TYPEDFIELDS) +""" +@kwdef struct Reference + "function `(x, dx) -> pf`" + pushforward = nothing + "function `(x, dy) -> pb`" + pullback = nothing + "function `x -> der`" + derivative = nothing + "function `x -> grad`" + gradient = nothing + "function `x -> jac`" + jacobian = nothing + "function `x -> der2`" + second_derivative = nothing + "function `(x, v) -> p`" + hvp = nothing + "function `x -> hess`" + hessian = nothing +end + +""" + Scenario{mutating} + +Store a testing scenario composed of a function and its input + output + tangents. + +# Fields + +$(TYPEDFIELDS) +""" +struct Scenario{mutating,F,X,Y,DX,DY,R} + "function" + f::F + "input" + x::X + "output" + y::Y + "pushforward seed" + dx::DX + "pullback seed" + dy::DY + "reference to compare against. It can be either an `ADTypes.AbstractADTypes` object or a [`Reference`](@ref) containing the correct operators associated with `f`" + ref::R +end + +function Scenario{mutating}( + f::F, x::X, y::Y, dx::DX, dy::DY, ref::R +) where {mutating,F,X,Y,DX,DY,R} + return Scenario{mutating,F,X,Y,DX,DY,R}(f, x, y, dx, dy, ref) +end + +function Base.string(scen::Scenario{mutating}) where {mutating} + return "$(string(scen.f)): $(typeof(scen.x)) -> $(typeof(scen.y))" +end + +is_mutating(::Scenario{mutating}) where {mutating} = mutating + +function change_ref(scen::Scenario{mutating}, new_ref::AbstractADType) where {mutating} + return Scenario{mutating}(scen.f, scen.x, scen.y, scen.dx, scen.dy, new_ref) +end + +## Scenario constructors + +function Scenario(f; x, y=nothing, ref=nothing) + if isnothing(y) + y = f(x) + dx = mysimilar_random(x) + dy = mysimilar_random(y) + return Scenario{false}(f, x, y, dx, dy, ref) + else + f! = f + f!(y, x) + dx = mysimilar_random(x) + dy = mysimilar_random(y) + return Scenario{true}(f!, x, y, dx, dy, ref) + end +end diff --git a/src/DifferentiationTest/scenarios_weird_arrays.jl b/src/DifferentiationTest/scenarios/weird_arrays.jl similarity index 64% rename from src/DifferentiationTest/scenarios_weird_arrays.jl rename to src/DifferentiationTest/scenarios/weird_arrays.jl index 7bfcbcd50..5982329ea 100644 --- a/src/DifferentiationTest/scenarios_weird_arrays.jl +++ b/src/DifferentiationTest/scenarios/weird_arrays.jl @@ -3,23 +3,26 @@ Create a vector of [`Scenario`](@ref)s involving weird types for testing differentiation. """ -function weird_array_scenarios(; static=true, component=true, gpu=true) +function weird_array_scenarios(; static=false, component=false, gpu=false) scenarios = Scenario[] if static ext = get_extension( DifferentiationInterface, :DifferentiationInterfaceStaticArraysExt ) - append!(scenarios, ext.static_scenarios_allocating()) + @assert !isnothing(ext) + append!(scenarios, ext.static_scenarios()) end if component ext = get_extension( DifferentiationInterface, :DifferentiationInterfaceComponentArraysExt ) - append!(scenarios, ext.component_scenarios_allocating()) + @assert !isnothing(ext) + append!(scenarios, ext.component_scenarios()) end if gpu ext = get_extension(DifferentiationInterface, :DifferentiationInterfaceJLArraysExt) - append!(scenarios, ext.gpu_scenarios_allocating()) + @assert !isnothing(ext) + append!(scenarios, ext.gpu_scenarios()) end return scenarios end diff --git a/src/DifferentiationTest/scenarios_default.jl b/src/DifferentiationTest/scenarios_default.jl deleted file mode 100644 index a0c49cb9f..000000000 --- a/src/DifferentiationTest/scenarios_default.jl +++ /dev/null @@ -1,116 +0,0 @@ -#= -Constraints on the scenarios: -- non-allocating whenever possible -- type-stable -- GPU-compatible (no scalar indexing) -- vary shapes to be tricky -=# - -const SCALING_VEC = Vector(1:12) -const SCALING_MAT = Matrix((1:3) .* transpose(1:4)) - -scalar_to_scalar(x::Number)::Number = sin(x) - -function scalar_to_vector(x::Number)::AbstractVector - return sin.(SCALING_VEC .* x) # output size 12 -end - -function scalar_to_vector!(y::AbstractVector, x::Number) - n = length(y) - y[1:(n ÷ 2)] .= sin.(x) - y[(n ÷ 2 + 1):n] .= sin.(2x) - return nothing -end - -function scalar_to_matrix(x::Number)::AbstractMatrix - return sin.(SCALING_MAT .* x) # output size (3, 4) -end - -function scalar_to_matrix!(y::AbstractMatrix, x::Number) - n, m = size(y) - y[1:(n ÷ 2), 1:(m ÷ 2)] .= sin.(x) - y[(n ÷ 2 + 1):n, 1:(m ÷ 2)] .= sin.(2x) - y[1:(n ÷ 2), ((m ÷ 2) + 1):m] .= cos.(x) - y[(n ÷ 2 + 1):n, ((m ÷ 2) + 1):m] .= cos.(2x) - return nothing -end - -vector_to_scalar(x::AbstractVector)::Number = sum(sin, x) -matrix_to_scalar(x::AbstractMatrix)::Number = sum(sin, x) - -vector_to_vector(x::AbstractVector)::AbstractVector = vcat(sin.(x), cos.(x)) - -function vector_to_vector!(y::AbstractVector, x::AbstractVector) - y[1:length(x)] .= sin.(x) - y[(length(x) + 1):(2length(x))] .= cos.(x) - return nothing -end - -vector_to_matrix(x::AbstractVector)::AbstractMatrix = hcat(sin.(x), cos.(x)) - -function vector_to_matrix!(y::AbstractMatrix, x::AbstractVector) - y[:, 1] .= sin.(x) - y[:, 2] .= cos.(x) - return nothing -end - -matrix_to_vector(x::AbstractMatrix)::AbstractVector = vcat(vec(sin.(x)), vec(cos.(x))) - -function matrix_to_vector!(y::AbstractVector, x::AbstractMatrix) - n = length(x) - y[1:n] .= sin.(getindex.(Ref(x), 1:n)) - y[(n + 1):(2n)] .= cos.(getindex.(Ref(x), 1:n)) - return nothing -end - -matrix_to_matrix(x::AbstractMatrix)::AbstractMatrix = hcat(vec(sin.(x)), vec(cos.(x))) - -function matrix_to_matrix!(y::AbstractMatrix, x::AbstractMatrix) - n = length(x) - y[:, 1] .= sin.(getindex.(Ref(x), 1:n)) - y[:, 2] .= cos.(getindex.(Ref(x), 1:n)) - return nothing -end - -function default_scenarios_allocating() - scenarios = [ - Scenario(scalar_to_scalar; x=2.0), - Scenario(scalar_to_vector; x=2.0), - Scenario(scalar_to_matrix; x=2.0), - Scenario(vector_to_scalar; x=Vector{Float64}(1:12)), - Scenario(matrix_to_scalar; x=Matrix{Float64}(reshape(1:12, 3, 4))), - Scenario(vector_to_vector; x=Vector{Float64}(1:12)), - Scenario(vector_to_matrix; x=Vector{Float64}(1:12)), - Scenario(matrix_to_vector; x=Matrix{Float64}(reshape(1:12, 3, 4))), - Scenario(matrix_to_matrix; x=Matrix{Float64}(reshape(1:12, 3, 4))), - ] - return scenarios -end - -function default_scenarios_mutating() - scenarios = [ - Scenario(scalar_to_vector!; x=2.0, y=zeros(Float64, length(SCALING_VEC))), - Scenario(scalar_to_matrix!; x=2.0, y=zeros(Float64, size(SCALING_MAT))), - Scenario(vector_to_vector!; x=Vector{Float64}(1:12), y=zeros(Float64, 24)), - Scenario(vector_to_matrix!; x=Vector{Float64}(1:12), y=zeros(Float64, 12, 2)), - Scenario( - matrix_to_vector!; x=Matrix{Float64}(reshape(1:12, 3, 4)), y=zeros(Float64, 24) - ), - Scenario( - matrix_to_matrix!; - x=Matrix{Float64}(reshape(1:12, 3, 4)), - y=zeros(Float64, 12, 2), - ), - ] - return scenarios -end - -""" - default_scenarios() - -Create a vector of [`Scenario`](@ref)s for testing differentiation. -""" -function default_scenarios() - scenarios = vcat(default_scenarios_allocating(), default_scenarios_mutating()) - return scenarios -end diff --git a/src/DifferentiationTest/scenarios_nested.jl b/src/DifferentiationTest/scenarios_nested.jl deleted file mode 100644 index 6d51da9ad..000000000 --- a/src/DifferentiationTest/scenarios_nested.jl +++ /dev/null @@ -1,50 +0,0 @@ -@kwdef struct Layer{W,B,A} - w::W - b::B - σ::A = nothing -end - -@functor Layer - -(l::Layer{<:Number,<:Number,<:Nothing})(x::Number) = l.w * x + l.b -(l::Layer{<:Number,<:Number})(x::Number) = l.σ(l.w * x + l.b) - -(l::Layer{<:AbstractMatrix,<:AbstractVector,<:Nothing})(x::AbstractVector) = l.w * x + l.b -(l::Layer{<:AbstractMatrix,<:AbstractVector})(x::AbstractVector) = l.σ.(l.w * x + l.b) - -call_layer(l::Layer{<:Number,<:Number}) = l(3.0) -call_layer(l::Layer{<:AbstractMatrix,<:AbstractVector}) = l(3 * ones(size(l.w, 2))) - -sum_call_layer(l) = sum(call_layer(l)) - -nested_norm(x::Number) = abs2(x) -nested_norm(x::AbstractArray) = sum(abs2, x) -nested_norm(x) = sum(nested_norm, fleaves(x)) - -function make_complicated(x::Number) - return (a=[2x, 3x], b=([exp(x);;],)) -end - -function make_complicated_with_immutables(x::Number) - return (a=[2x, 3x], b=([exp(x);;],), c=sin(x)) -end - -function nested_scenarios(; immutables=true) - scenarios_without_immutables = [ - Scenario(make_complicated, 2.0), - Scenario(nested_norm; x=make_complicated(2.0)), - Scenario(sum_call_layer; x=Layer(; w=rand(2, 3), b=rand(2))), - Scenario(sum_call_layer; x=Layer(; w=rand(2, 3), b=rand(2), σ=tanh)), - ] - scenarios_with_immutables = [ - Scenario(make_complicated_with_immutables, 2.0), - Scenario(nested_norm; x=make_complicated_with_immutables(2.0)), - Scenario(call_layer; x=Layer(; w=2.0, b=4.0)), - Scenario(call_layer; x=Layer(; w=2.0, b=4.0, σ=tanh)), - ] - if immutables - return vcat(scenarios_with_immutables, scenarios_without_immutables) - else - return scenarios_without_immutables - end -end diff --git a/src/DifferentiationTest/test_error_free.jl b/src/DifferentiationTest/test_error_free.jl deleted file mode 100644 index dcc38b0c7..000000000 --- a/src/DifferentiationTest/test_error_free.jl +++ /dev/null @@ -1,143 +0,0 @@ - -function test_error_free(ba::AbstractADType, ::typeof(pushforward), scen::Scenario{false}) - (; f, x, dx, dy) = deepcopy(scen) - extras = prepare_pushforward(f, ba, x) - dy_in = mysimilar(dy) - - @test (value_and_pushforward!!(f, dy_in, ba, x, dx, extras); true) - @test (value_and_pushforward(f, ba, x, dx, extras); true) - @test (pushforward!!(f, dy_in, ba, x, dx, extras); true) - @test (pushforward(f, ba, x, dx, extras); true) - return nothing -end - -function test_error_free(ba::AbstractADType, ::typeof(pushforward), scen::Scenario{true}) - (; f, x, y, dx, dy) = deepcopy(scen) - f! = f - extras = prepare_pushforward(f!, ba, y, x) - y_in = mysimilar(y) - dy_in = mysimilar(dy) - @test (value_and_pushforward!!(f!, y_in, dy_in, ba, x, dx, extras); true) - return nothing -end - -## Pullback - -function test_error_free(ba::AbstractADType, ::typeof(pullback), scen::Scenario{false}) - (; f, x, dx, dy) = deepcopy(scen) - extras = prepare_pullback(f, ba, x) - dx_in = mysimilar(dx) - - @test (value_and_pullback!!(f, dx_in, ba, x, dy, extras); true) - @test (value_and_pullback(f, ba, x, dy, extras); true) - @test (pullback!!(f, dx_in, ba, x, dy, extras); true) - @test (pullback(f, ba, x, dy, extras); true) - return nothing -end - -function test_error_free(ba::AbstractADType, ::typeof(pullback), scen::Scenario{true}) - (; f, x, y, dx, dy) = deepcopy(scen) - f! = f - extras = prepare_pullback(f!, ba, y, x) - y_in = mysimilar(y) - dx_in = mysimilar(dx) - - @test (value_and_pullback!!(f!, y_in, dx_in, ba, x, dy, extras); true) - return nothing -end - -## Derivative - -function test_error_free(ba::AbstractADType, ::typeof(derivative), scen::Scenario{false}) - (; f, x, dy) = deepcopy(scen) - extras = prepare_derivative(f, ba, x) - der_in = mysimilar(dy) - - @test (value_and_derivative!!(f, der_in, ba, x, extras); true) - @test (value_and_derivative(f, ba, x, extras); true) - @test (derivative!!(f, der_in, ba, x, extras); true) - @test (derivative(f, ba, x, extras); true) - return nothing -end - -function test_error_free(ba::AbstractADType, ::typeof(derivative), scen::Scenario{true}) - (; f, x, y, dy) = deepcopy(scen) - f! = f - extras = prepare_derivative(f!, ba, y, x) - y_in = mysimilar(y) - der_in = mysimilar(dy) - - @test (value_and_derivative!!(f!, y_in, der_in, ba, x, extras); true) - return nothing -end - -## Gradient - -function test_error_free(ba::AbstractADType, ::typeof(gradient), scen::Scenario{false}) - (; f, x, dx) = deepcopy(scen) - extras = prepare_gradient(f, ba, x) - grad_in = mysimilar(dx) - - @test (value_and_gradient!!(f, grad_in, ba, x, extras); true) - @test (value_and_gradient(f, ba, x, extras); true) - @test (gradient!!(f, grad_in, ba, x, extras); true) - @test (gradient(f, ba, x, extras); true) - return nothing -end - -## Jacobian - -function test_error_free(ba::AbstractADType, ::typeof(jacobian), scen::Scenario{false}) - (; f, x, y) = deepcopy(scen) - extras = prepare_jacobian(f, ba, x) - jac_in = similar(x, length(y), length(x)) - - @test (value_and_jacobian!!(f, jac_in, ba, x, extras); true) - @test (value_and_jacobian(f, ba, x, extras); true) - @test (jacobian!!(f, jac_in, ba, x, extras); true) - @test (jacobian(f, ba, x, extras); true) - return nothing -end - -function test_error_free(ba::AbstractADType, ::typeof(jacobian), scen::Scenario{true}) - (; f, x, y) = deepcopy(scen) - f! = f - extras = prepare_jacobian(f!, ba, y, x) - y_in = mysimilar(y) - jac_in = similar(x, length(y), length(x)) - - @test (value_and_jacobian!!(f!, y_in, jac_in, ba, x, extras); true) - return nothing -end - -## Second derivative - -function test_error_free( - ba::AbstractADType, ::typeof(second_derivative), scen::Scenario{false}; -) - (; f, x) = deepcopy(scen) - extras = prepare_second_derivative(f, ba, x) - - @test (second_derivative(f, ba, x, extras); true) - return nothing -end - -## HVP - -function test_error_free(ba::AbstractADType, ::typeof(hvp), scen::Scenario{false}) - (; f, x, dx) = deepcopy(scen) - extras = prepare_hvp(f, ba, x) - - @test (hvp(f, ba, x, dx, extras); true) - return nothing -end - -## Hessian - -function test_error_free(ba::AbstractADType, ::typeof(hessian), scen::Scenario{false}) - (; f, x) = deepcopy(scen) - extras = prepare_hessian(f, ba, x) - - @test (hessian(f, ba, x, extras); true) - return nothing -end diff --git a/src/DifferentiationTest/benchmark.jl b/src/DifferentiationTest/tests/benchmark.jl similarity index 100% rename from src/DifferentiationTest/benchmark.jl rename to src/DifferentiationTest/tests/benchmark.jl diff --git a/src/DifferentiationTest/test_call_count.jl b/src/DifferentiationTest/tests/call_count.jl similarity index 98% rename from src/DifferentiationTest/test_call_count.jl rename to src/DifferentiationTest/tests/call_count.jl index d559f39fd..d3e7d49a5 100644 --- a/src/DifferentiationTest/test_call_count.jl +++ b/src/DifferentiationTest/tests/call_count.jl @@ -124,7 +124,7 @@ function test_call_count(ba::AbstractADType, ::typeof(jacobian), scen::Scenario{ extras = prepare_jacobian(CallCounter(f), ba, y, x) cc! = CallCounter(f) y_in = mysimilar(y) - jac_in = similar(y, length(y), length(x)) + jac_in = Matrix{eltype(y)}(undef, length(y), length(x)) value_and_jacobian!!(cc!, y_in, jac_in, ba, x, extras) if mode(ba) == AbstractForwardMode @test cc!.count[] <= 1 + length(x) diff --git a/src/DifferentiationTest/tests/correctness.jl b/src/DifferentiationTest/tests/correctness.jl new file mode 100644 index 000000000..d43dced2d --- /dev/null +++ b/src/DifferentiationTest/tests/correctness.jl @@ -0,0 +1,356 @@ +## No overwrite + +function test_scen_intact(new_scen, scen) + let (≈)(x, y) = myisapprox(x, y; rtol=0) + @testset "Scenario intact" begin + @test new_scen.x ≈ scen.x + @test new_scen.y ≈ scen.y + @test new_scen.dx ≈ scen.dx + @test new_scen.dy ≈ scen.dy + end + end +end + +## Pushforward + +function test_correctness( + ba::AbstractADType, ::typeof(pushforward), scen::Scenario{false}; rtol +) + (; f, x, y, dx, dy, ref) = new_scen = deepcopy(scen) + dy_true = if ref isa AbstractADType + pushforward(f, ref, x, dx) + else + ref.pushforward(x, dx) + end + + y1, dy1 = value_and_pushforward(f, ba, x, dx) + y2, dy2 = value_and_pushforward!!(f, mysimilar(dy), ba, x, dx) + + dy3 = pushforward(f, ba, x, dx) + dy4 = pushforward!!(f, mysimilar(dy), ba, x, dx) + + let (≈)(x, y) = myisapprox(x, y; rtol) + @testset "Primal value" begin + @test y1 ≈ y + @test y2 ≈ y + end + @testset "Tangent value" begin + @test dy1 ≈ dy_true + @test dy2 ≈ dy_true + @test dy3 ≈ dy_true + @test dy4 ≈ dy_true + end + end + test_scen_intact(new_scen, scen) + return nothing +end + +function test_correctness( + ba::AbstractADType, ::typeof(pushforward), scen::Scenario{true}; rtol +) + (; f, x, y, dx, dy, ref) = new_scen = deepcopy(scen) + f! = f + dy_true = if ref isa AbstractADType + last(value_and_pushforward!!(f!, mysimilar(y), mysimilar(dy), ref, x, dx)) + else + ref.pushforward(x, dx) + end + + y10 = mysimilar(y) + y1, dy1 = value_and_pushforward!!(f!, y10, mysimilar(dy), ba, x, dx) + + let (≈)(x, y) = myisapprox(x, y; rtol) + @testset "Primal value" begin + @test y10 ≈ y + @test y1 ≈ y + end + @testset "Tangent value" begin + @test dy1 ≈ dy_true + end + end + test_scen_intact(new_scen, scen) + return nothing +end + +## Pullback + +function test_correctness( + ba::AbstractADType, ::typeof(pullback), scen::Scenario{false}; rtol +) + (; f, x, y, dx, dy, ref) = new_scen = deepcopy(scen) + dx_true = if ref isa AbstractADType + pullback(f, ref, x, dy) + else + ref.pullback(x, dy) + end + + y1, dx1 = value_and_pullback(f, ba, x, dy) + y2, dx2 = value_and_pullback!!(f, mysimilar(dx), ba, x, dy) + + dx3 = pullback(f, ba, x, dy) + dx4 = pullback!!(f, mysimilar(dx), ba, x, dy) + + let (≈)(x, y) = myisapprox(x, y; rtol) + @testset "Primal value" begin + @test y1 ≈ y + @test y2 ≈ y + end + @testset "Cotangent value" begin + @test dx1 ≈ dx_true + @test dx2 ≈ dx_true + @test dx3 ≈ dx_true + @test dx4 ≈ dx_true + end + end + test_scen_intact(new_scen, scen) + return nothing +end + +function test_correctness( + ba::AbstractADType, ::typeof(pullback), scen::Scenario{true}; rtol +) + (; f, x, y, dx, dy, ref) = new_scen = deepcopy(scen) + f! = f + dx_true = if ref isa AbstractADType + last(value_and_pullback!!(f, mysimilar(y), mysimilar(dx), ref, x, dy)) + else + ref.pullback(x, dy) + end + + y10 = mysimilar(y) + y1, dx1 = value_and_pullback!!(f!, y10, mysimilar(dx), ba, x, dy) + + let (≈)(x, y) = myisapprox(x, y; rtol) + @testset "Primal value" begin + @test y10 ≈ y + @test y1 ≈ y + end + @testset "Cotangent value" begin + @test dx1 ≈ dx_true + end + end + test_scen_intact(new_scen, scen) + return nothing +end + +## Derivative + +function test_correctness( + ba::AbstractADType, ::typeof(derivative), scen::Scenario{false}; rtol +) + (; f, x, y, dx, dy, ref) = new_scen = deepcopy(scen) + der_true = if ref isa AbstractADType + derivative(f, ref, x) + else + ref.derivative(x) + end + + y1, der1 = value_and_derivative(f, ba, x) + y2, der2 = value_and_derivative!!(f, mysimilar(dy), ba, x) + + der3 = derivative(f, ba, x) + der4 = derivative!!(f, mysimilar(dy), ba, x) + + let (≈)(x, y) = myisapprox(x, y; rtol) + @testset "Primal value" begin + @test y1 ≈ y + @test y2 ≈ y + end + @testset "Derivative value" begin + @test der1 ≈ der_true + @test der2 ≈ der_true + @test der3 ≈ der_true + @test der4 ≈ der_true + end + end + test_scen_intact(new_scen, scen) + return nothing +end + +function test_correctness( + ba::AbstractADType, ::typeof(derivative), scen::Scenario{true}; rtol +) + (; f, x, y, dx, dy, ref) = new_scen = deepcopy(scen) + f! = f + der_true = if ref isa AbstractADType + last(value_and_derivative!!(f!, mysimilar(y), mysimilar(dy), ref, x)) + else + ref.derivative(x) + end + + y10 = mysimilar(y) + y1, der1 = value_and_derivative!!(f!, y10, mysimilar(dy), ba, x) + + let (≈)(x, y) = myisapprox(x, y; rtol) + @testset "Primal value" begin + @test y10 ≈ y + @test y1 ≈ y + end + @testset "Derivative value" begin + @test der1 ≈ der_true + end + end + test_scen_intact(new_scen, scen) + return nothing +end + +## Gradient + +function test_correctness( + ba::AbstractADType, ::typeof(gradient), scen::Scenario{false}; rtol +) + (; f, x, y, dx, dy, ref) = new_scen = deepcopy(scen) + grad_true = if ref isa AbstractADType + gradient(f, ref, x) + else + ref.gradient(x) + end + + y1, grad1 = value_and_gradient(f, ba, x) + y2, grad2 = value_and_gradient!!(f, mysimilar(dx), ba, x) + + grad3 = gradient(f, ba, x) + grad4 = gradient!!(f, mysimilar(dx), ba, x) + + let (≈)(x, y) = myisapprox(x, y; rtol) + @testset "Primal value" begin + @test y1 ≈ y + @test y2 ≈ y + end + @testset "Gradient value" begin + @test grad1 ≈ grad_true + @test grad2 ≈ grad_true + @test grad3 ≈ grad_true + @test grad4 ≈ grad_true + end + end + test_scen_intact(new_scen, scen) + return nothing +end + +## Jacobian + +function test_correctness( + ba::AbstractADType, ::typeof(jacobian), scen::Scenario{false}; rtol +) + (; f, x, y, ref) = new_scen = deepcopy(scen) + jac_true = if ref isa AbstractADType + jacobian(f, ref, x) + else + ref.jacobian(x) + end + + y1, jac1 = value_and_jacobian(f, ba, x) + y2, jac2 = value_and_jacobian!!(f, mysimilar(jac_true), ba, x) + + jac3 = jacobian(f, ba, x) + jac4 = jacobian!!(f, mysimilar(jac_true), ba, x) + + let (≈)(x, y) = myisapprox(x, y; rtol) + @testset "Primal value" begin + @test y1 ≈ y + @test y2 ≈ y + end + @testset "Jacobian value" begin + @test jac1 ≈ jac_true + @test jac2 ≈ jac_true + @test jac3 ≈ jac_true + @test jac4 ≈ jac_true + end + end + test_scen_intact(new_scen, scen) + return nothing +end + +function test_correctness( + ba::AbstractADType, ::typeof(jacobian), scen::Scenario{true}; rtol +) + (; f, x, y, dy, ref) = new_scen = deepcopy(scen) + f! = f + jac_shape = Matrix{eltype(y)}(undef, length(y), length(x)) + jac_true = if ref isa AbstractADType + last(value_and_jacobian!!(f!, mysimilar(y), mysimilar(jac_shape), ref, x)) + else + ref.jacobian(x) + end + + y10 = mysimilar(y) + y1, jac1 = value_and_jacobian!!(f!, y10, mysimilar(jac_true), ba, x) + + let (≈)(x, y) = myisapprox(x, y; rtol) + @testset "Primal value" begin + @test y10 ≈ y + @test y1 ≈ y + end + @testset "Jacobian value" begin + @test jac1 ≈ jac_true + end + end + test_scen_intact(new_scen, scen) + return nothing +end + +## Second derivative + +function test_correctness( + ba::AbstractADType, ::typeof(second_derivative), scen::Scenario; rtol +) + (; f, x, 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) + + let (≈)(x, y) = myisapprox(x, y; rtol) + @testset "Second derivative value" begin + @test der21 ≈ der2_true + end + end + test_scen_intact(new_scen, scen) + return nothing +end + +## Hessian-vector product + +function test_correctness(ba::AbstractADType, ::typeof(hvp), scen::Scenario; rtol) + (; f, x, dx, ref) = new_scen = deepcopy(scen) + hvp_true = if ref isa AbstractADType + hvp(f, ref, x, dx) + else + ref.hvp(x, dx) + end + + hvp1 = hvp(f, ba, x, dx) + + let (≈)(x, y) = myisapprox(x, y; rtol) + @testset "HVP value" begin + @test hvp1 ≈ hvp_true + end + end + test_scen_intact(new_scen, scen) + return nothing +end + +## Hessian + +function test_correctness(ba::AbstractADType, ::typeof(hessian), scen::Scenario; rtol) + (; f, x, y, ref) = new_scen = deepcopy(scen) + hess_true = if ref isa AbstractADType + hessian(f, ref, x) + else + ref.hessian(x) + end + + hess1 = hessian(f, ba, x) + + let (≈)(x, y) = myisapprox(x, y; rtol) + @testset "Hessian value" begin + @test hess1 ≈ hess_true + end + end + test_scen_intact(new_scen, scen) + return nothing +end diff --git a/src/DifferentiationTest/test_differentiation.jl b/src/DifferentiationTest/tests/test.jl similarity index 68% rename from src/DifferentiationTest/test_differentiation.jl rename to src/DifferentiationTest/tests/test.jl index a82c09468..72c9da8e7 100644 --- a/src/DifferentiationTest/test_differentiation.jl +++ b/src/DifferentiationTest/tests/test.jl @@ -52,8 +52,8 @@ end Cross-test a list of `backends` for a list of `operators` on a list of `scenarios`, running a variety of different tests. -- If `benchmark` is `false`, this returns a `TestSet` object. -- If `benchmark` is `true`, this returns a [`BenchmarkData`](@ref) object, which is easy to turn into a `DataFrame`. +- If `benchmark` is `false`, this runs the tests and returns `nothing`. +- If `benchmark` is `true`, this runs the tests and returns a [`BenchmarkData`](@ref) object, which is easy to turn into a `DataFrame`. # Default arguments @@ -64,10 +64,9 @@ Cross-test a list of `backends` for a list of `operators` on a list of `scenario Testing: -- `correctness=true`: whether to compare the differentiation results with those given by ForwardDiff.jl -- `error_free=false`: whether to run it once and see if it errors +- `correctness=true`: whether to compare the differentiation results with the theoretical values specified in each scenario. If a backend object like `correctness=AutoForwardDiff()` is passed instead of a boolean, the results will be compared using that reference backend as the ground truth. - `call_count=false`: whether to check that the function is called the right number of times -- `type_stability=false`: whether to check type stability with JET.jl (`@test_opt`) +- `type_stability=false`: whether to check type stability with JET.jl (thanks to `@test_opt`) - `benchmark=false`: whether to run and return a benchmark suite with Chairmarks.jl - `allocations=false`: whether to check that the benchmarks are allocation-free - `detailed=false`: whether to print a detailed test set (by scenario) or condensed test set (by operator) @@ -76,19 +75,22 @@ Filtering: - `input_type=Any`: restrict scenario inputs to subtypes of this - `output_type=Any`: restrict scenario outputs to subtypes of this -- `first_order=true`: consider first order operators -- `second_order=true`: consider second order operators - `allocating=true`: consider operators for allocating functions - `mutating=true`: consider operators for mutating functions +- `first_order=true`: consider first order operators +- `second_order=true`: consider second order operators - `excluded=Symbol[]`: list of excluded operators + +Options: + +- `rtol=1e-3`: precision for correctness testing (when comparing to the reference outputs) """ function test_differentiation( backends::Vector{<:AbstractADType}, operators::Vector{<:Function}=all_operators(), scenarios::Vector{<:Scenario}=default_scenarios(); # testing - correctness::Bool=true, - error_free::Bool=false, + correctness::Union{Bool,AbstractADType}=true, type_stability::Bool=false, call_count::Bool=false, benchmark::Bool=false, @@ -97,11 +99,13 @@ function test_differentiation( # filtering input_type::Type=Any, output_type::Type=Any, - first_order=true, - second_order=true, allocating=true, mutating=true, + first_order=true, + second_order=true, excluded::Vector{<:Function}=Function[], + # options + rtol=1e-3, ) operators = filter_operators(operators; first_order, second_order, excluded) scenarios = filter_scenarios(scenarios; input_type, output_type, allocating, mutating) @@ -110,23 +114,12 @@ function test_differentiation( title = "Differentiation tests -" * - (correctness ? " correctness" : "") * - (error_free ? " errors" : "") * + (correctness != false ? " correctness" : "") * (call_count ? " calls" : "") * (type_stability ? " types" : "") * (benchmark ? " benchmark" : "") * (allocations ? " allocations" : "") - correctness_ext = if correctness - ext = get_extension( - DifferentiationInterface, :DifferentiationInterfaceCorrectnessTestExt - ) - @assert !isnothing(ext) - ext - else - nothing - end - jet_ext = if type_stability ext = get_extension(DifferentiationInterface, :DifferentiationInterfaceJETExt) @assert !isnothing(ext) @@ -143,38 +136,37 @@ function test_differentiation( nothing end - test_set = @testset verbose = true "$title" begin - @testset verbose = true "$(backend_string(backend))" for backend in backends - @testset verbose = detailed "$op" for op in operators - @testset "$scen" for scen in filter(scenarios) do scen - compatible(backend, op, scen) - end - if correctness - @testset "Correctness" begin - correctness_ext.test_correctness(backend, op, scen) - end - end - if error_free - @testset "Error-free" begin - test_error_free(backend, op, scen) + @testset verbose = detailed "$(backend_string(backend))" for backend in backends + @testset verbose = detailed "$op" for op in operators + @testset "$scen" for scen in filter(scenarios) do scen + compatible(backend, op, scen) + end + if correctness != false + @testset "Correctness" begin + if correctness isa AbstractADType + test_correctness( + backend, op, change_ref(scen, correctness); rtol + ) + else + test_correctness(backend, op, scen; rtol) end end - if call_count - @testset "Call count" begin - test_call_count(backend, op, scen) - end + end + if call_count + @testset "Call count" begin + test_call_count(backend, op, scen) end - if type_stability - @testset "Type stability" begin - jet_ext.test_jet(backend, op, scen) - end + end + if type_stability + @testset "Type stability" begin + jet_ext.test_jet(backend, op, scen) end - if benchmark || allocations - @testset "Allocations" begin - chairmarks_ext.run_benchmark!( - benchmark_data, backend, op, scen; allocations=allocations - ) - end + end + if benchmark || allocations + @testset "Allocations" begin + chairmarks_ext.run_benchmark!( + benchmark_data, backend, op, scen; allocations=allocations + ) end end end @@ -184,7 +176,7 @@ function test_differentiation( if benchmark return benchmark_data else - return test_set + return nothing end end diff --git a/src/DifferentiationTest/compatibility.jl b/src/DifferentiationTest/utils/compatibility.jl similarity index 100% rename from src/DifferentiationTest/compatibility.jl rename to src/DifferentiationTest/utils/compatibility.jl diff --git a/src/DifferentiationTest/printing.jl b/src/DifferentiationTest/utils/printing.jl similarity index 100% rename from src/DifferentiationTest/printing.jl rename to src/DifferentiationTest/utils/printing.jl diff --git a/src/DifferentiationTest/zero.jl b/src/DifferentiationTest/utils/zero.jl similarity index 100% rename from src/DifferentiationTest/zero.jl rename to src/DifferentiationTest/utils/zero.jl diff --git a/src/backends.jl b/src/backends.jl index 66228e6bc..ee77d26b3 100644 --- a/src/backends.jl +++ b/src/backends.jl @@ -11,7 +11,7 @@ function check_available(backend::AbstractADType) value_and_gradient(abs2, backend, 2.0) return true catch exception - # @warn "Backend $backend not available" exception + @warn "Backend $backend not available" exception if exception isa MethodError return false else @@ -35,7 +35,7 @@ function check_mutation(backend::AbstractADType) y, jac = value_and_jacobian!!(square!, [0.0], [0.0;;], backend, [3.0]) return isapprox(y, [9.0]; rtol=1e-3) && isapprox(jac, [6.0;;]; rtol=1e-3) catch exception - # @warn "Backend $backend does not support mutation" exception + @warn "Backend $backend does not support mutation" exception return false end end @@ -56,7 +56,7 @@ function check_hessian(backend::AbstractADType) hess = hessian(sqnorm, backend, x) return isapprox(hess, [2.0;;]; rtol=1e-3) catch exception - # @warn "Backend $backend does not support hessian" exception + @warn "Backend $backend does not support hessian" exception return false end end diff --git a/test/Project.toml b/test/Project.toml index 83ec71d1a..19f2ed850 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -6,6 +6,7 @@ Chairmarks = "0ca39b1e-fe0b-4e98-acfc-b1656634c4de" ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66" DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" Diffractor = "9f5e2b26-1114-432f-b630-d3fe2085c51c" +Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" FastDifferentiation = "eb9bf01b-bf85-4b60-bf87-ee5de06c00be" FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41" diff --git a/test/chainrules_reverse.jl b/test/chainrules_reverse.jl deleted file mode 100644 index fea24663b..000000000 --- a/test/chainrules_reverse.jl +++ /dev/null @@ -1,9 +0,0 @@ -include("test_imports.jl") - -using Zygote: ZygoteRuleConfig - -@test check_available(AutoChainRules(ZygoteRuleConfig())) -@test !check_mutation(AutoChainRules(ZygoteRuleConfig())) -@test_broken !check_hessian(AutoChainRules(ZygoteRuleConfig())) - -test_differentiation(AutoChainRules(ZygoteRuleConfig()); second_order=false); diff --git a/test/diffractor.jl b/test/diffractor.jl deleted file mode 100644 index 962f221d1..000000000 --- a/test/diffractor.jl +++ /dev/null @@ -1,9 +0,0 @@ -include("test_imports.jl") - -using Diffractor: Diffractor - -@test check_available(AutoDiffractor()) -@test !check_mutation(AutoDiffractor()) -@test_broken check_hessian(AutoDiffractor()) - -test_differentiation(AutoDiffractor(); second_order=false); diff --git a/test/enzyme_forward.jl b/test/enzyme_forward.jl deleted file mode 100644 index 931d13229..000000000 --- a/test/enzyme_forward.jl +++ /dev/null @@ -1,8 +0,0 @@ -include("test_imports.jl") - -using Enzyme: Enzyme - -@test check_available(AutoEnzyme(Enzyme.Forward)) -@test check_mutation(AutoEnzyme(Enzyme.Forward)) - -test_differentiation(AutoEnzyme(Enzyme.Forward); type_stability=true, second_order=false); diff --git a/test/enzyme_reverse.jl b/test/enzyme_reverse.jl deleted file mode 100644 index 00263c887..000000000 --- a/test/enzyme_reverse.jl +++ /dev/null @@ -1,8 +0,0 @@ -include("test_imports.jl") - -using Enzyme: Enzyme - -@test check_available(AutoEnzyme(Enzyme.Reverse)) -@test check_mutation(AutoEnzyme(Enzyme.Reverse)) - -test_differentiation(AutoEnzyme(Enzyme.Reverse); second_order=false); diff --git a/test/fastdifferentiation.jl b/test/fastdifferentiation.jl deleted file mode 100644 index 62964c5b4..000000000 --- a/test/fastdifferentiation.jl +++ /dev/null @@ -1,15 +0,0 @@ -include("test_imports.jl") - -using DifferentiationInterface: AutoFastDifferentiation -using FastDifferentiation: FastDifferentiation - -@test check_available(AutoFastDifferentiation()) -@test !check_mutation(AutoFastDifferentiation()) -@test_broken !check_hessian(AutoFastDifferentiation()) - -test_differentiation( - AutoFastDifferentiation(); - input_type=Union{Number,AbstractVector}, - output_type=Union{Number,AbstractVector}, - second_order=false, -); diff --git a/test/finitediff.jl b/test/finitediff.jl deleted file mode 100644 index 068c88604..000000000 --- a/test/finitediff.jl +++ /dev/null @@ -1,9 +0,0 @@ -include("test_imports.jl") - -using FiniteDiff: FiniteDiff - -@test check_available(AutoFiniteDiff()) -@test check_mutation(AutoFiniteDiff()) -@test_broken check_hessian(AutoFiniteDiff()) - -test_differentiation(AutoFiniteDiff(); second_order=false); diff --git a/test/finitedifferences.jl b/test/finitedifferences.jl deleted file mode 100644 index 15d529512..000000000 --- a/test/finitedifferences.jl +++ /dev/null @@ -1,9 +0,0 @@ -include("test_imports.jl") - -using FiniteDifferences: FiniteDifferences, central_fdm - -@test check_available(AutoFiniteDifferences(central_fdm(5, 1))) -@test !check_mutation(AutoFiniteDifferences(central_fdm(5, 1))) -@test_broken !check_hessian(AutoFiniteDifferences(central_fdm(5, 1))) - -test_differentiation(AutoFiniteDifferences(central_fdm(5, 1)); second_order=false); diff --git a/test/first_order.jl b/test/first_order.jl new file mode 100644 index 000000000..74db8b7b6 --- /dev/null +++ b/test/first_order.jl @@ -0,0 +1,36 @@ +include("test_imports.jl") + +using ChainRulesCore: ChainRulesCore +using Diffractor: Diffractor +using Enzyme: Enzyme +using FastDifferentiation: FastDifferentiation +using FiniteDiff: FiniteDiff +using FiniteDifferences: FiniteDifferences +using ForwardDiff: ForwardDiff +using PolyesterForwardDiff: PolyesterForwardDiff +using ReverseDiff: ReverseDiff +using Tracker: Tracker +using Zygote: Zygote + +## + +all_backends = [ + AutoChainRules(Zygote.ZygoteRuleConfig()), + AutoDiffractor(), + AutoEnzyme(Enzyme.Forward), + AutoEnzyme(Enzyme.Reverse), + AutoFastDifferentiation(), + AutoFiniteDiff(), + AutoFiniteDifferences(FiniteDifferences.central_fdm(3, 1)), + AutoForwardDiff(), + AutoPolyesterForwardDiff(; chunksize=2), + AutoReverseDiff(), + AutoTracker(), + AutoZygote(), +] + +for backend in all_backends + @test check_available(backend) +end + +test_differentiation(all_backends; second_order=false); diff --git a/test/forwarddiff.jl b/test/forwarddiff.jl deleted file mode 100644 index 24c8f4af5..000000000 --- a/test/forwarddiff.jl +++ /dev/null @@ -1,16 +0,0 @@ -include("test_imports.jl") - -using ForwardDiff: ForwardDiff - -@test check_available(AutoForwardDiff()) -@test check_mutation(AutoForwardDiff()) -@test check_hessian(AutoForwardDiff()) - -test_differentiation(AutoForwardDiff(; chunksize=2); type_stability=true); - -test_differentiation( - AutoForwardDiff(; chunksize=2), - all_operators(), - weird_array_scenarios(; static=true, component=true, gpu=false); - excluded=[jacobian], -); diff --git a/test/nested.jl b/test/nested.jl index d0ed64c8a..2e3f441a4 100644 --- a/test/nested.jl +++ b/test/nested.jl @@ -4,17 +4,10 @@ using ForwardDiff: ForwardDiff using Tracker: Tracker using Zygote: Zygote +## Reverse mode + test_differentiation(AutoZygote(), [gradient], nested_scenarios(); detailed=true); -test_differentiation( - AutoTracker(), [gradient], nested_scenarios(; immutables=false); detailed=true -); +## Forward mode -test_differentiation( - AutoForwardDiff(), - [derivative], - nested_scenarios(); - detailed=true, - correctness=false, - error_free=true, -); +test_differentiation(AutoForwardDiff(), [derivative], nested_scenarios(); detailed=true); diff --git a/test/nobackend.jl b/test/nobackend.jl deleted file mode 100644 index f06598193..000000000 --- a/test/nobackend.jl +++ /dev/null @@ -1,9 +0,0 @@ -include("test_imports.jl") - -using ADTypes: ADTypes - -struct AutoNothingForward <: ADTypes.AbstractForwardMode end -struct AutoNothingReverse <: ADTypes.AbstractReverseMode end - -@test !check_available(AutoNothingForward()) -@test !check_available(AutoNothingReverse()) diff --git a/test/polyesterforwarddiff.jl b/test/polyesterforwarddiff.jl deleted file mode 100644 index c295a6a73..000000000 --- a/test/polyesterforwarddiff.jl +++ /dev/null @@ -1,14 +0,0 @@ -include("test_imports.jl") - -using PolyesterForwardDiff: PolyesterForwardDiff - -@test check_available(AutoPolyesterForwardDiff(; chunksize=2)) -@test check_mutation(AutoPolyesterForwardDiff(; chunksize=2)) -@test check_hessian(AutoPolyesterForwardDiff(; chunksize=2)) - -test_differentiation( - AutoPolyesterForwardDiff(; chunksize=2); - input_type=Union{Number,AbstractVector}, - output_type=Union{Number,AbstractVector}, - type_stability=false, -); diff --git a/test/reversediff.jl b/test/reversediff.jl deleted file mode 100644 index d207a11db..000000000 --- a/test/reversediff.jl +++ /dev/null @@ -1,9 +0,0 @@ -include("test_imports.jl") - -using ReverseDiff: ReverseDiff - -@test check_available(AutoReverseDiff()) -@test check_mutation(AutoReverseDiff()) -@test check_hessian(AutoReverseDiff()) - -test_differentiation(AutoReverseDiff()); diff --git a/test/runtests.jl b/test/runtests.jl index 79897a534..2c04fa32a 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -5,87 +5,43 @@ include("test_imports.jl") @testset verbose = true "DifferentiationInterface.jl" begin @testset verbose = true "Formal tests" begin @testset "Aqua" begin - @info "Running Aqua.jl tests..." - @time Aqua.test_all(DifferentiationInterface; ambiguities=false) + Aqua.test_all(DifferentiationInterface; ambiguities=false) end @testset "JuliaFormatter" begin - @info "Running JuliaFormatter test..." - @time @test JuliaFormatter.format( + @test JuliaFormatter.format( DifferentiationInterface; verbose=false, overwrite=false ) end @testset "JET" begin - @info "Running JET tests..." - @time JET.test_package(DifferentiationInterface; target_defined_modules=true) + JET.test_package(DifferentiationInterface; target_defined_modules=true) end end - @testset verbose = true "Trivial backends" begin - @info "Running no backend tests..." - @testset "No backend" begin - @time include("nobackend.jl") - end - @info "Running zero backend tests..." - @testset "Zero backend" begin - @time include("zero.jl") - end + Documenter.doctest(DifferentiationInterface) + + @testset "Zero backend" begin + include("zero.jl") end @testset verbose = true "First order" begin - @testset "ChainRules (reverse)" begin - @info "Running ChainRules (reverse) tests..." - @time include("chainrules_reverse.jl") - end - @testset "Diffractor (forward)" begin - @info "Running Diffractor (forward) tests..." - @time include("diffractor.jl") - end - @testset "Enzyme (forward)" begin - @info "Running Enzyme (forward) tests..." - @time include("enzyme_forward.jl") - end - @testset "Enzyme (reverse)" begin - @info "Running Enzyme (reverse) tests..." - @time include("enzyme_reverse.jl") - end - @testset "FastDifferentiation" begin - @info "Running FastDifferentiation tests..." - @time include("fastdifferentiation.jl") - end - @testset "FiniteDiff" begin - @info "Running FiniteDiff tests..." - @time include("finitediff.jl") - end - @testset "FiniteDifferences" begin - @info "Running FiniteDifferences tests..." - @time include("finitedifferences.jl") - end - @testset "ForwardDiff" begin - @info "Running ForwardDiff tests..." - @time include("forwarddiff.jl") - end - @testset "PolyesterForwardDiff" begin - @info "Running PolyesterForwardDiff tests..." - @time include("polyesterforwarddiff.jl") - end - @testset "ReverseDiff" begin - @info "Running ReverseDiff tests..." - @time include("reversediff.jl") - end - @testset "Tracker" begin - @info "Running Tracker tests..." - @time include("tracker.jl") - end - @testset "Zygote" begin - @info "Running Zygote tests..." - @time include("zygote.jl") - end + include("first_order.jl") + end + + @testset verbose = true "Second order" begin + include("second_order.jl") end - include("second_order.jl") + @testset verbose = true "Bonus round" begin + @testset "Type stability" begin + include("type_stability.jl") + end - @testset "Weird types" begin - @info "Testing nested objects..." - @time include("nested.jl") + @testset "Weird arrays" begin + include("weird_arrays.jl") + end + + @testset "Nested types" begin + include("nested.jl") + end end end; diff --git a/test/second_order.jl b/test/second_order.jl index 7f3b60c19..6521f41f1 100644 --- a/test/second_order.jl +++ b/test/second_order.jl @@ -1,25 +1,20 @@ -using DifferentiationInterface.DifferentiationTest: backend_string +include("test_imports.jl") -using FiniteDiff: FiniteDiff -using ForwardDiff: ForwardDiff using Enzyme: Enzyme -using Zygote: Zygote +using ForwardDiff: ForwardDiff +using ReverseDiff: ReverseDiff + +second_order_backends = [AutoForwardDiff(), AutoReverseDiff()] + +second_order_mixed_backends = [ + SecondOrder(AutoEnzyme(Enzyme.Forward), AutoForwardDiff()), + SecondOrder(AutoForwardDiff(), AutoEnzyme(Enzyme.Forward)), + SecondOrder(AutoForwardDiff(), AutoZygote()), +] -SECOND_ORDER_BACKENDS = Dict( - "forward/forward" => [ - SecondOrder(AutoEnzyme(Enzyme.Forward), AutoForwardDiff()), - SecondOrder(AutoForwardDiff(), AutoEnzyme(Enzyme.Forward)), - ], - "forward/reverse" => [SecondOrder(AutoForwardDiff(), AutoZygote())], - "reverse/forward" => [], -) +for backend in vcat(second_order_backends, second_order_mixed_backends) + @test check_hessian(backend) +end -@testset verbose = true "Second order" begin - @testset verbose = true "$second_order_mode" for (second_order_mode, backends) in - pairs(SECOND_ORDER_BACKENDS) - @info "Testing $second_order_mode..." - @time @testset "$(backend_string(backend))" for backend in backends - test_differentiation(backend; first_order=false, type_stability=false) - end - end -end; +test_differentiation(second_order_backends; first_order=false, second_order=true); +test_differentiation(second_order_mixed_backends; first_order=false, second_order=true); diff --git a/test/taped.jl b/test/taped.jl deleted file mode 100644 index 51b865adf..000000000 --- a/test/taped.jl +++ /dev/null @@ -1,13 +0,0 @@ -using Pkg -Pkg.add(; url="https://github.com/withbayes/Taped.jl/") - -include("test_imports.jl") - -using DifferentiationInterface: AutoTaped -using Taped: Taped - -@test check_available(AutoTaped()) -@test !check_mutation(AutoTaped()) -@test !check_hessian(AutoTaped()) - -test_differentiation(AutoTaped(); second_order=false, excluded=[jacobian]); diff --git a/test/test_imports.jl b/test/test_imports.jl index d0f4dbedc..0faa7e7cf 100644 --- a/test/test_imports.jl +++ b/test/test_imports.jl @@ -3,8 +3,10 @@ using ADTypes using DifferentiationInterface using DifferentiationInterface.DifferentiationTest +using DifferentiationInterface.DifferentiationTest: backend_string using Aqua: Aqua +using Documenter: Documenter using JET: JET using JuliaFormatter: JuliaFormatter using Test @@ -15,4 +17,3 @@ using ForwardDiff: ForwardDiff using ComponentArrays using JLArrays using StaticArrays -using Zygote: Zygote diff --git a/test/tracker.jl b/test/tracker.jl deleted file mode 100644 index a77ce9577..000000000 --- a/test/tracker.jl +++ /dev/null @@ -1,11 +0,0 @@ -include("test_imports.jl") - -using Tracker: Tracker - -@test check_available(AutoTracker()) -@test !check_mutation(AutoTracker()) -@test !check_hessian(AutoTracker()) - -test_differentiation( - AutoTracker(); output_type=Union{Number,AbstractVector}, second_order=false -); diff --git a/test/type_stability.jl b/test/type_stability.jl new file mode 100644 index 000000000..f58631b56 --- /dev/null +++ b/test/type_stability.jl @@ -0,0 +1,12 @@ +include("test_imports.jl") + +using Enzyme: Enzyme +using ForwardDiff: ForwardDiff + +type_stable_backends = [ + AutoForwardDiff(), AutoEnzyme(Enzyme.Forward), AutoEnzyme(Enzyme.Reverse) +] + +test_differentiation( + type_stable_backends; correctness=false, type_stability=true, second_order=false +); diff --git a/test/weird_arrays.jl b/test/weird_arrays.jl new file mode 100644 index 000000000..bdecd25b2 --- /dev/null +++ b/test/weird_arrays.jl @@ -0,0 +1,21 @@ +include("test_imports.jl") + +using ForwardDiff: ForwardDiff +using Zygote: Zygote + +test_differentiation( + AutoForwardDiff(; chunksize=2), + all_operators(), + # ForwardDiff access individual indices + weird_array_scenarios(; static=true, component=true, gpu=false); + # jacobian is super long for some reason + excluded=[jacobian], + second_order=false, +); + +test_differentiation( + AutoZygote(), + all_operators(), + weird_array_scenarios(; static=true, component=true, gpu=true); + second_order=false, +); diff --git a/test/zero.jl b/test/zero.jl index 5590a52c8..887322bb5 100644 --- a/test/zero.jl +++ b/test/zero.jl @@ -5,31 +5,31 @@ using DifferentiationInterface.DifferentiationTest: AutoZeroForward, AutoZeroRev @test check_available(AutoZeroForward()) @test check_available(AutoZeroReverse()) -## Error-free & type-stability +## Correctness (vs oneself) + type-stability -test_differentiation( - [AutoZeroForward(), AutoZeroReverse()]; - correctness=false, - error_free=true, - type_stability=true, -); +for backend in [AutoZeroForward(), AutoZeroReverse()] + test_differentiation( + backend, + all_operators(), + default_scenarios(); + correctness=backend, + type_stability=true, + ) +end -test_differentiation( - [AutoZeroForward(), AutoZeroReverse()], - all_operators(), - weird_array_scenarios(; static=true, component=false, gpu=true); - correctness=false, - error_free=true, -); - -test_differentiation( - [AutoZeroForward(), AutoZeroReverse()], - all_operators(), - weird_array_scenarios(; static=false, component=true, gpu=false); - correctness=false, - error_free=true, - excluded=[hessian], -); +for backend in [ + SecondOrder(AutoZeroForward(), AutoZeroReverse()), + SecondOrder(AutoZeroReverse(), AutoZeroForward()), +] + test_differentiation( + backend, + all_operators(), + default_scenarios(); + correctness=backend, + type_stability=true, + first_order=false, + ) +end ## Call count @@ -55,3 +55,27 @@ data = test_differentiation( ); df = DataFrames.DataFrame(pairs(data)...) + +## Weird arrays + +for backend in [AutoZeroForward(), AutoZeroReverse()] + test_differentiation( + backend, all_operators(), weird_array_scenarios(; gpu=true); correctness=backend + ) + # copyto!(col, col) fails on static arrays + test_differentiation( + backend, + all_operators(), + weird_array_scenarios(; static=true); + correctness=backend, + excluded=[jacobian], + ) + # stack fails on component vectors + test_differentiation( + backend, + all_operators(), + weird_array_scenarios(; component=true); + correctness=backend, + excluded=[hessian], + ) +end diff --git a/test/zygote.jl b/test/zygote.jl deleted file mode 100644 index 078fb8b4b..000000000 --- a/test/zygote.jl +++ /dev/null @@ -1,16 +0,0 @@ -include("test_imports.jl") - -using Zygote: Zygote - -@test check_available(AutoZygote()) -@test !check_mutation(AutoZygote()) -@test_skip !check_hessian(AutoZygote()) - -test_differentiation(AutoZygote(); second_order=false); - -test_differentiation( - AutoZygote(), - all_operators(), - weird_array_scenarios(; static=true, component=true, gpu=true); - second_order=false, -);