From 0a29640655366442ca78adc2f8a093c249957bae Mon Sep 17 00:00:00 2001 From: Adrian Hill Date: Fri, 22 Mar 2024 22:29:17 +0100 Subject: [PATCH] Add operator traits for testing (#90) --- .JuliaFormatter.toml | 7 +- .../DifferentiationInterfaceChairmarksExt.jl | 200 ++++++++++-------- .../DifferentiationInterfaceForwardDiffExt.jl | 16 ++ .../test_correctness.jl | 114 +++------- .../DifferentiationInterfaceJETExt.jl | 114 ++++------ .../DifferentiationTest.jl | 8 +- src/DifferentiationTest/call_count.jl | 128 ++++------- src/DifferentiationTest/default_scenarios.jl | 60 +++--- src/DifferentiationTest/operator_traits.jl | 115 ++++++++++ src/DifferentiationTest/scenario.jl | 40 ++-- src/DifferentiationTest/test_operators.jl | 73 ++++--- test/runtests.jl | 15 ++ 12 files changed, 454 insertions(+), 436 deletions(-) create mode 100644 src/DifferentiationTest/operator_traits.jl diff --git a/.JuliaFormatter.toml b/.JuliaFormatter.toml index c7439503e..cc91afe1d 100644 --- a/.JuliaFormatter.toml +++ b/.JuliaFormatter.toml @@ -1 +1,6 @@ -style = "blue" \ No newline at end of file +style = "blue" +align_assignment = true +align_struct_field = true +align_conditional = true +align_pair_arrow = true +align_matrix = true diff --git a/ext/DifferentiationInterfaceChairmarksExt/DifferentiationInterfaceChairmarksExt.jl b/ext/DifferentiationInterfaceChairmarksExt/DifferentiationInterfaceChairmarksExt.jl index 895f7bfd0..d8c756c01 100644 --- a/ext/DifferentiationInterfaceChairmarksExt/DifferentiationInterfaceChairmarksExt.jl +++ b/ext/DifferentiationInterfaceChairmarksExt/DifferentiationInterfaceChairmarksExt.jl @@ -18,11 +18,27 @@ using DifferentiationInterface: supports_hvp using DifferentiationInterface.DifferentiationTest import DifferentiationInterface.DifferentiationTest as DT +using DifferentiationInterface.DifferentiationTest: + AbstractOperator, + PushforwardAllocating, + PushforwardMutating, + PullbackAllocating, + PullbackMutating, + MultiderivativeAllocating, + MultiderivativeMutating, + GradientAllocating, + JacobianAllocating, + JacobianMutating, + DerivativeAllocating, + SecondDerivativeAllocating, + HessianAllocating, + HessianVectorProductAllocating, + compatible_scenarios using Test function DT.run_benchmark( backends::Vector{<:AbstractADType}, - operators::Vector{Symbol}, + operators::Vector{<:AbstractOperator}, scenarios::Vector{<:Scenario}; allocations=false, ) @@ -30,72 +46,8 @@ function DT.run_benchmark( @testset verbose = true "Allocations" begin @testset verbose = true "$(backend_string(backend))" for backend in backends @testset "$op" for op in operators - if op == :pushforward_allocating - @testset "$s" for s in allocating(scenarios) - benchmark_pushforward_allocating!(data, backend, s; allocations) - end - elseif op == :pushforward_mutating - @testset "$s" for s in mutating(scenarios) - benchmark_pushforward_mutating!(data, backend, s; allocations) - end - - elseif op == :pullback_allocating - @testset "$s" for s in allocating(scenarios) - benchmark_pullback_allocating!(data, backend, s; allocations) - end - elseif op == :pullback_mutating - @testset "$s" for s in mutating(scenarios) - benchmark_pullback_mutating!(data, backend, s; allocations) - end - - elseif op == :derivative_allocating - @testset "$s" for s in allocating(scalar_scalar(scenarios)) - benchmark_derivative_allocating!(data, backend, s; allocations) - end - - elseif op == :multiderivative_allocating - @testset "$s" for s in allocating(scalar_array(scenarios)) - benchmark_multiderivative_allocating!(data, backend, s; allocations) - end - elseif op == :multiderivative_mutating - @testset "$s" for s in mutating(scalar_array(scenarios)) - benchmark_multiderivative_mutating!(data, backend, s; allocations) - end - - elseif op == :gradient_allocating - @testset "$s" for s in allocating(array_scalar(scenarios)) - benchmark_gradient_allocating!(data, backend, s; allocations) - end - - elseif op == :jacobian_allocating - @testset "$s" for s in allocating(array_array(scenarios)) - benchmark_jacobian_allocating!(data, backend, s; allocations) - end - elseif op == :jacobian_mutating - @testset "$s" for s in mutating(array_array(scenarios)) - benchmark_jacobian_mutating!(data, backend, s; allocations) - end - - elseif op == :second_derivative_allocating - @testset "$s" for s in allocating(scalar_scalar(scenarios)) - benchmark_second_derivative_allocating!( - data, backend, s; allocations - ) - end - - elseif op == :hessian_vector_product_allocating - @testset "$s" for s in allocating(array_scalar(scenarios)) - benchmark_hessian_vector_product_allocating!( - data, backend, s; allocations - ) - end - elseif op == :hessian_allocating - @testset "$s" for s in allocating(array_scalar(scenarios)) - benchmark_hessian_allocating!(data, backend, s; allocations) - end - - else - throw(ArgumentError("Invalid operator to benchmark: `:$op`")) + @testset "$s" for s in compatible_scenarios(op, scenarios) + benchmark!(op, data, backend, s; allocations) end end end @@ -105,8 +57,12 @@ end ## Pushforward -function benchmark_pushforward_allocating!( - data::BenchmarkData, ba::AbstractADType, scen::Scenario; allocations::Bool +function benchmark!( + ::PushforwardAllocating, + data::BenchmarkData, + ba::AbstractADType, + scen::Scenario; + allocations::Bool, ) Bool(supports_pushforward(ba)) || return nothing (; f, x, dx, dy) = deepcopy(scen) @@ -123,8 +79,12 @@ function benchmark_pushforward_allocating!( return nothing end -function benchmark_pushforward_mutating!( - data::BenchmarkData, ba::AbstractADType, scen::Scenario; allocations::Bool +function benchmark!( + ::PushforwardMutating, + data::BenchmarkData, + ba::AbstractADType, + scen::Scenario; + allocations::Bool, ) Bool(supports_pushforward(ba)) || return nothing Bool(supports_mutation(ba)) || return nothing @@ -143,8 +103,12 @@ end ## Pullback -function benchmark_pullback_allocating!( - data::BenchmarkData, ba::AbstractADType, scen::Scenario; allocations::Bool +function benchmark!( + ::PullbackAllocating, + data::BenchmarkData, + ba::AbstractADType, + scen::Scenario; + allocations::Bool, ) Bool(supports_pullback(ba)) || return nothing (; f, x, dx, dy) = deepcopy(scen) @@ -160,8 +124,12 @@ function benchmark_pullback_allocating!( return nothing end -function benchmark_pullback_mutating!( - data::BenchmarkData, ba::AbstractADType, scen::Scenario; allocations::Bool +function benchmark!( + ::PullbackMutating, + data::BenchmarkData, + ba::AbstractADType, + scen::Scenario; + allocations::Bool, ) Bool(supports_pullback(ba)) || return nothing Bool(supports_mutation(ba)) || return nothing @@ -178,8 +146,12 @@ end ## Derivative -function benchmark_derivative_allocating!( - data::BenchmarkData, ba::AbstractADType, scen::Scenario; allocations::Bool +function benchmark!( + ::DerivativeAllocating, + data::BenchmarkData, + ba::AbstractADType, + scen::Scenario; + allocations::Bool, ) (; f, x) = deepcopy(scen) extras = prepare_derivative(ba, f, x) @@ -193,8 +165,12 @@ end ## Multiderivative -function benchmark_multiderivative_allocating!( - data::BenchmarkData, ba::AbstractADType, scen::Scenario; allocations::Bool +function benchmark!( + ::MultiderivativeAllocating, + data::BenchmarkData, + ba::AbstractADType, + scen::Scenario; + allocations::Bool, ) (; f, x, dy) = deepcopy(scen) extras = prepare_multiderivative(ba, f, x) @@ -204,8 +180,12 @@ function benchmark_multiderivative_allocating!( return nothing end -function benchmark_multiderivative_mutating!( - data::BenchmarkData, ba::AbstractADType, scen::Scenario; allocations::Bool +function benchmark!( + ::MultiderivativeMutating, + data::BenchmarkData, + ba::AbstractADType, + scen::Scenario; + allocations::Bool, ) Bool(supports_mutation(ba)) || return nothing (; f, x, y, dy) = deepcopy(scen) @@ -223,8 +203,12 @@ end ## Gradient -function benchmark_gradient_allocating!( - data::BenchmarkData, ba::AbstractADType, scen::Scenario; allocations::Bool +function benchmark!( + ::GradientAllocating, + data::BenchmarkData, + ba::AbstractADType, + scen::Scenario; + allocations::Bool, ) (; f, x, dx) = deepcopy(scen) extras = prepare_gradient(ba, f, x) @@ -241,8 +225,12 @@ end ## Jacobian -function benchmark_jacobian_allocating!( - data::BenchmarkData, ba::AbstractADType, scen::Scenario; allocations::Bool +function benchmark!( + ::JacobianAllocating, + data::BenchmarkData, + ba::AbstractADType, + scen::Scenario; + allocations::Bool, ) (; f, x, y) = deepcopy(scen) jac_template = zeros(eltype(y), length(y), length(x)) @@ -253,8 +241,12 @@ function benchmark_jacobian_allocating!( return nothing end -function benchmark_jacobian_mutating!( - data::BenchmarkData, ba::AbstractADType, scen::Scenario; allocations::Bool +function benchmark!( + ::JacobianMutating, + data::BenchmarkData, + ba::AbstractADType, + scen::Scenario; + allocations::Bool, ) Bool(supports_mutation(ba)) || return nothing (; f, x, y) = deepcopy(scen) @@ -273,8 +265,12 @@ end ## Second derivative -function benchmark_second_derivative_allocating!( - data::BenchmarkData, ba::AbstractADType, scen::Scenario; allocations::Bool +function benchmark!( + ::SecondDerivativeAllocating, + data::BenchmarkData, + ba::AbstractADType, + scen::Scenario; + allocations::Bool, ) (; f, x) = deepcopy(scen) extras = prepare_second_derivative(ba, f, x) @@ -288,8 +284,12 @@ end ## Hessian-vector product -function benchmark_hessian_vector_product_allocating!( - data::BenchmarkData, ba::AbstractADType, scen::Scenario; allocations::Bool +function benchmark!( + ::HessianVectorProductAllocating, + data::BenchmarkData, + ba::AbstractADType, + scen::Scenario; + allocations::Bool, ) Bool(supports_hvp(ba)) || return nothing (; f, x, dx) = deepcopy(scen) @@ -309,8 +309,12 @@ end ## Hessian -function benchmark_hessian_allocating!( - data::BenchmarkData, ba::AbstractADType, scen::Scenario; allocations::Bool +function benchmark!( + ::HessianAllocating, + data::BenchmarkData, + ba::AbstractADType, + scen::Scenario; + allocations::Bool, ) (; f, x, y, dx) = deepcopy(scen) extras = prepare_hessian(ba, f, x) @@ -328,4 +332,14 @@ function benchmark_hessian_allocating!( return nothing end +function benchmark!( + op::AbstractOperator, + data::BenchmarkData, + ba::AbstractADType, + scen::Scenario; + allocations::Bool, +) + throw(ArgumentError("Invalid operator to test: $op")) end + +end # module diff --git a/ext/DifferentiationInterfaceForwardDiffExt/DifferentiationInterfaceForwardDiffExt.jl b/ext/DifferentiationInterfaceForwardDiffExt/DifferentiationInterfaceForwardDiffExt.jl index 0a13bade5..723084f91 100644 --- a/ext/DifferentiationInterfaceForwardDiffExt/DifferentiationInterfaceForwardDiffExt.jl +++ b/ext/DifferentiationInterfaceForwardDiffExt/DifferentiationInterfaceForwardDiffExt.jl @@ -18,6 +18,22 @@ using DifferentiationInterface: import DifferentiationInterface as DI using DifferentiationInterface.DifferentiationTest import DifferentiationInterface.DifferentiationTest as DT +using DifferentiationInterface.DifferentiationTest: + AbstractOperator, + PushforwardAllocating, + PushforwardMutating, + PullbackAllocating, + PullbackMutating, + MultiderivativeAllocating, + MultiderivativeMutating, + GradientAllocating, + JacobianAllocating, + JacobianMutating, + DerivativeAllocating, + SecondDerivativeAllocating, + HessianAllocating, + HessianVectorProductAllocating, + compatible_scenarios using DiffResults: DiffResults using DocStringExtensions using ForwardDiff: diff --git a/ext/DifferentiationInterfaceForwardDiffExt/test_correctness.jl b/ext/DifferentiationInterfaceForwardDiffExt/test_correctness.jl index 741f96198..e3bdc9faf 100644 --- a/ext/DifferentiationInterfaceForwardDiffExt/test_correctness.jl +++ b/ext/DifferentiationInterfaceForwardDiffExt/test_correctness.jl @@ -3,74 +3,14 @@ function DT.test_correctness( backends::Vector{<:AbstractADType}, - operators::Vector{Symbol}, + operators::Vector{<:AbstractOperator}, scenarios::Vector{<:Scenario}; ) @testset verbose = true "Correctness" begin @testset verbose = true "$(backend_string(backend))" for backend in backends @testset "$op" for op in operators - if op == :pushforward_allocating - @testset "$s" for s in allocating(scenarios) - test_correctness_pushforward_allocating(backend, s) - end - elseif op == :pushforward_mutating - @testset "$s" for s in mutating(scenarios) - test_correctness_pushforward_mutating(backend, s) - end - - elseif op == :pullback_allocating - @testset "$s" for s in allocating(scenarios) - test_correctness_pullback_allocating(backend, s) - end - elseif op == :pullback_mutating - @testset "$s" for s in mutating(scenarios) - test_correctness_pullback_mutating(backend, s) - end - - elseif op == :derivative_allocating - @testset "$s" for s in allocating(scalar_scalar(scenarios)) - test_correctness_derivative_allocating(backend, s) - end - - elseif op == :multiderivative_allocating - @testset "$s" for s in allocating(scalar_array(scenarios)) - test_correctness_multiderivative_allocating(backend, s) - end - elseif op == :multiderivative_mutating - @testset "$s" for s in mutating(scalar_array(scenarios)) - test_correctness_multiderivative_mutating(backend, s) - end - - elseif op == :gradient_allocating - @testset "$s" for s in allocating(array_scalar(scenarios)) - test_correctness_gradient_allocating(backend, s) - end - - elseif op == :jacobian_allocating - @testset "$s" for s in allocating(array_array(scenarios)) - test_correctness_jacobian_allocating(backend, s) - end - elseif op == :jacobian_mutating - @testset "$s" for s in mutating(array_array(scenarios)) - test_correctness_jacobian_mutating(backend, s) - end - - elseif op == :second_derivative_allocating - @testset "$s" for s in allocating(scalar_scalar(scenarios)) - test_correctness_second_derivative_allocating(backend, s) - end - - elseif op == :hessian_vector_product_allocating - @testset "$s" for s in allocating(array_scalar(scenarios)) - test_correctness_hessian_vector_product_allocating(backend, s) - end - elseif op == :hessian_allocating - @testset "$s" for s in allocating(array_scalar(scenarios)) - test_correctness_hessian_allocating(backend, s) - end - - else - throw(ArgumentError("Invalid operator to test: `:$op`")) + @testset "$s" for s in compatible_scenarios(op, scenarios) + test_correctness(op, backend, s) end end end @@ -79,7 +19,7 @@ end ## Pushforward -function test_correctness_pushforward_allocating(ba::AbstractADType, scen::Scenario) +function test_correctness(::PushforwardAllocating, ba::AbstractADType, scen::Scenario) Bool(supports_pushforward(ba)) || return nothing (; f, x, y, dx) = deepcopy(scen) dy_true = true_pushforward(f, x, y, dx; mutating=false) @@ -110,7 +50,7 @@ function test_correctness_pushforward_allocating(ba::AbstractADType, scen::Scena end end -function test_correctness_pushforward_mutating(ba::AbstractADType, scen::Scenario) +function test_correctness(::PushforwardMutating, ba::AbstractADType, scen::Scenario) Bool(supports_pushforward(ba)) || return nothing Bool(supports_mutation(ba)) || return nothing (; f, x, y, dx) = deepcopy(scen) @@ -137,7 +77,7 @@ end ## Pullback -function test_correctness_pullback_allocating(ba::AbstractADType, scen::Scenario) +function test_correctness(::PullbackAllocating, ba::AbstractADType, scen::Scenario) Bool(supports_pullback(ba)) || return nothing (; f, x, y, dy) = deepcopy(scen) dx_true = true_pullback(f, x, y, dy; mutating=false) @@ -168,7 +108,7 @@ function test_correctness_pullback_allocating(ba::AbstractADType, scen::Scenario end end -function test_correctness_pullback_mutating(ba::AbstractADType, scen::Scenario) +function test_correctness(::PullbackMutating, ba::AbstractADType, scen::Scenario) Bool(supports_pullback(ba)) || return nothing Bool(supports_mutation(ba)) || return nothing (; f, x, y, dy) = deepcopy(scen) @@ -197,7 +137,7 @@ end ## Derivative -function test_correctness_derivative_allocating(ba::AbstractADType, scen::Scenario) +function test_correctness(::DerivativeAllocating, ba::AbstractADType, scen::Scenario) (; f, x, y) = deepcopy(scen) der_true = ForwardDiff.derivative(f, x) @@ -216,7 +156,7 @@ end ## Multiderivative -function test_correctness_multiderivative_allocating(ba::AbstractADType, scen::Scenario) +function test_correctness(::MultiderivativeAllocating, ba::AbstractADType, scen::Scenario) (; f, x, y) = deepcopy(scen) multider_true = ForwardDiff.derivative(f, x) @@ -244,7 +184,7 @@ function test_correctness_multiderivative_allocating(ba::AbstractADType, scen::S end end -function test_correctness_multiderivative_mutating(ba::AbstractADType, scen::Scenario) +function test_correctness(::MultiderivativeMutating, ba::AbstractADType, scen::Scenario) Bool(supports_mutation(ba)) || return nothing (; f, x, y) = deepcopy(scen) f! = f @@ -270,7 +210,7 @@ end ## Gradient -function test_correctness_gradient_allocating(ba::AbstractADType, scen::Scenario) +function test_correctness(::GradientAllocating, ba::AbstractADType, scen::Scenario) (; f, x, y) = deepcopy(scen) grad_true = ForwardDiff.gradient(f, x) @@ -300,7 +240,7 @@ end ## Jacobian -function test_correctness_jacobian_allocating(ba::AbstractADType, scen::Scenario) +function test_correctness(::JacobianAllocating, ba::AbstractADType, scen::Scenario) (; f, x, y) = deepcopy(scen) jac_true = ForwardDiff.jacobian(f, x) @@ -328,7 +268,7 @@ function test_correctness_jacobian_allocating(ba::AbstractADType, scen::Scenario end end -function test_correctness_jacobian_mutating(ba::AbstractADType, scen::Scenario) +function test_correctness(::JacobianMutating, ba::AbstractADType, scen::Scenario) Bool(supports_mutation(ba)) || return nothing (; f, x, y) = deepcopy(scen) f! = f @@ -354,7 +294,7 @@ end ## Second derivative -function test_correctness_second_derivative_allocating(ba::AbstractADType, scen::Scenario) +function test_correctness(::SecondDerivativeAllocating, ba::AbstractADType, scen::Scenario) (; f, x, y) = deepcopy(scen) der_true = ForwardDiff.derivative(f, x) derder_true = ForwardDiff.derivative(x) do z @@ -379,8 +319,8 @@ end ## Hessian-vector product -function test_correctness_hessian_vector_product_allocating( - ba::AbstractADType, scen::Scenario +function test_correctness( + ::HessianVectorProductAllocating, ba::AbstractADType, scen::Scenario ) Bool(supports_hvp(ba)) || return nothing (; f, x, dx) = deepcopy(scen) @@ -420,7 +360,7 @@ end ## Hessian -function test_correctness_hessian_allocating(ba::AbstractADType, scen::Scenario) +function test_correctness(::HessianAllocating, ba::AbstractADType, scen::Scenario) (; f, x, y) = deepcopy(scen) grad_true = ForwardDiff.gradient(f, x) hess_true = ForwardDiff.hessian(f, x) @@ -458,13 +398,17 @@ function test_correctness_hessian_allocating(ba::AbstractADType, scen::Scenario) end end +function test_correctness(op::AbstractOperator, ba::AbstractADType, scen::Scenario) + throw(ArgumentError("Invalid operator to test: $op")) +end + ## Utils -function true_pushforward(f, x::Number, y::Number, dx; mutating) +function true_pushforward(f::F, x::Number, y::Number, dx; mutating) where {F} return ForwardDiff.derivative(f, x) * dx end -function true_pushforward(f, x::Number, y::AbstractArray, dx; mutating) +function true_pushforward(f::F, x::Number, y::AbstractArray, dx; mutating) where {F} if mutating return ForwardDiff.derivative(f, deepcopy(y), x) .* dx else @@ -472,11 +416,11 @@ function true_pushforward(f, x::Number, y::AbstractArray, dx; mutating) end end -function true_pushforward(f, x::AbstractArray, y::Number, dx; mutating) +function true_pushforward(f::F, x::AbstractArray, y::Number, dx; mutating) where {F} return dot(ForwardDiff.gradient(f, x), dx) end -function true_pushforward(f, x::AbstractArray, y::AbstractArray, dx; mutating) +function true_pushforward(f::F, x::AbstractArray, y::AbstractArray, dx; mutating) where {F} if mutating return reshape(ForwardDiff.jacobian(f, deepcopy(y), x) * vec(dx), size(y)) else @@ -484,11 +428,11 @@ function true_pushforward(f, x::AbstractArray, y::AbstractArray, dx; mutating) end end -function true_pullback(f, x::Number, y::Number, dy; mutating) +function true_pullback(f::F, x::Number, y::Number, dy; mutating) where {F} return ForwardDiff.derivative(f, x) * dy end -function true_pullback(f, x::Number, y::AbstractArray, dy; mutating) +function true_pullback(f::F, x::Number, y::AbstractArray, dy; mutating) where {F} if mutating return dot(ForwardDiff.derivative(f, deepcopy(y), x), dy) else @@ -496,11 +440,11 @@ function true_pullback(f, x::Number, y::AbstractArray, dy; mutating) end end -function true_pullback(f, x::AbstractArray, y::Number, dy; mutating) +function true_pullback(f::F, x::AbstractArray, y::Number, dy; mutating) where {F} return ForwardDiff.gradient(f, x) .* dy end -function true_pullback(f, x::AbstractArray, y::AbstractArray, dy; mutating) +function true_pullback(f::F, x::AbstractArray, y::AbstractArray, dy; mutating) where {F} if mutating return reshape( transpose(ForwardDiff.jacobian(f, deepcopy(y), x)) * vec(dy), size(x) diff --git a/ext/DifferentiationInterfaceJETExt/DifferentiationInterfaceJETExt.jl b/ext/DifferentiationInterfaceJETExt/DifferentiationInterfaceJETExt.jl index fe9e09493..6fa597916 100644 --- a/ext/DifferentiationInterfaceJETExt/DifferentiationInterfaceJETExt.jl +++ b/ext/DifferentiationInterfaceJETExt/DifferentiationInterfaceJETExt.jl @@ -17,6 +17,22 @@ using DifferentiationInterface: supports_hvp using DifferentiationInterface.DifferentiationTest import DifferentiationInterface.DifferentiationTest as DT +using DifferentiationInterface.DifferentiationTest: + AbstractOperator, + PushforwardAllocating, + PushforwardMutating, + PullbackAllocating, + PullbackMutating, + MultiderivativeAllocating, + MultiderivativeMutating, + GradientAllocating, + JacobianAllocating, + JacobianMutating, + DerivativeAllocating, + SecondDerivativeAllocating, + HessianAllocating, + HessianVectorProductAllocating, + compatible_scenarios using JET: @test_call, @test_opt using LinearAlgebra: LinearAlgebra using Test @@ -25,83 +41,23 @@ using Test function DT.test_type_stability( backends::Vector{<:AbstractADType}, - operators::Vector{Symbol}, + operators::Vector{<:AbstractOperator}, scenarios::Vector{<:Scenario}; ) @testset verbose = true "Type stability" begin @testset verbose = true "$(backend_string(backend))" for backend in backends @testset "$op" for op in operators - if op == :pushforward_allocating - @testset "$s" for s in allocating(scenarios) - test_type_pushforward_allocating(backend, s) - end - elseif op == :pushforward_mutating - @testset "$s" for s in mutating(scenarios) - test_type_pushforward_mutating(backend, s) - end - - elseif op == :pullback_allocating - @testset "$s" for s in allocating(scenarios) - test_type_pullback_allocating(backend, s) - end - elseif op == :pullback_mutating - @testset "$s" for s in mutating(scenarios) - test_type_pullback_mutating(backend, s) - end - - elseif op == :derivative_allocating - @testset "$s" for s in allocating(scalar_scalar(scenarios)) - test_type_derivative_allocating(backend, s) - end - - elseif op == :multiderivative_allocating - @testset "$s" for s in allocating(scalar_array(scenarios)) - test_type_multiderivative_allocating(backend, s) - end - elseif op == :multiderivative_mutating - @testset "$s" for s in mutating(scalar_array(scenarios)) - test_type_multiderivative_mutating(backend, s) - end - - elseif op == :gradient_allocating - @testset "$s" for s in allocating(array_scalar(scenarios)) - test_type_gradient_allocating(backend, s) - end - - elseif op == :jacobian_allocating - @testset "$s" for s in allocating(array_array(scenarios)) - test_type_jacobian_allocating(backend, s) - end - elseif op == :jacobian_mutating - @testset "$s" for s in mutating(array_array(scenarios)) - test_type_jacobian_mutating(backend, s) - end - - elseif op == :second_derivative_allocating - @testset "$s" for s in allocating(scalar_scalar(scenarios)) - test_type_second_derivative_allocating(backend, s) - end - - elseif op == :hessian_vector_product_allocating - @testset "$s" for s in allocating(array_scalar(scenarios)) - test_type_hessian_vector_product_allocating(backend, s) - end - elseif op == :hessian_allocating - @testset "$s" for s in allocating(array_scalar(scenarios)) - test_type_hessian_allocating(backend, s) - end - - else - throw(ArgumentError("Invalid operator to test: `:$op`")) + @testset "$s" for s in compatible_scenarios(op, scenarios) + test_type(op, backend, s) end end end end end -## Pushforward +## Pushforward -function test_type_pushforward_allocating(ba::AbstractADType, scen::Scenario) +function test_type(::PushforwardAllocating, ba::AbstractADType, scen::Scenario) Bool(supports_pushforward(ba)) || return nothing (; f, x, dx, dy) = deepcopy(scen) dy_in = zero(dy) @@ -119,7 +75,7 @@ function test_type_pushforward_allocating(ba::AbstractADType, scen::Scenario) @test_opt pushforward(ba, f, x, dx) end -function test_type_pushforward_mutating(ba::AbstractADType, scen::Scenario) +function test_type(::PushforwardMutating, ba::AbstractADType, scen::Scenario) Bool(supports_pushforward(ba)) || return nothing Bool(supports_mutation(ba)) || return nothing (; f, x, y, dx, dy) = deepcopy(scen) @@ -133,7 +89,7 @@ end ## Pullback -function test_type_pullback_allocating(ba::AbstractADType, scen::Scenario) +function test_type(::PullbackAllocating, ba::AbstractADType, scen::Scenario) Bool(supports_pullback(ba)) || return nothing (; f, x, dx, dy) = deepcopy(scen) dx_in = zero(dx) @@ -151,7 +107,7 @@ function test_type_pullback_allocating(ba::AbstractADType, scen::Scenario) @test_opt pullback(ba, f, x, dy) end -function test_type_pullback_mutating(ba::AbstractADType, scen::Scenario) +function test_type(::PullbackMutating, ba::AbstractADType, scen::Scenario) Bool(supports_pullback(ba)) || return nothing Bool(supports_mutation(ba)) || return nothing (; f, x, y, dx, dy) = deepcopy(scen) @@ -165,7 +121,7 @@ end ## Derivative -function test_type_derivative_allocating(ba::AbstractADType, scen::Scenario) +function test_type(::DerivativeAllocating, ba::AbstractADType, scen::Scenario) (; f, x) = deepcopy(scen) @test_call value_and_derivative(ba, f, x) @@ -177,7 +133,7 @@ end ## Multiderivative -function test_type_multiderivative_allocating(ba::AbstractADType, scen::Scenario) +function test_type(::MultiderivativeAllocating, ba::AbstractADType, scen::Scenario) (; f, x, dy) = deepcopy(scen) multider_in = zero(dy) @@ -194,7 +150,7 @@ function test_type_multiderivative_allocating(ba::AbstractADType, scen::Scenario @test_opt multiderivative(ba, f, x) end -function test_type_multiderivative_mutating(ba::AbstractADType, scen::Scenario) +function test_type(::MultiderivativeMutating, ba::AbstractADType, scen::Scenario) Bool(supports_mutation(ba)) || return nothing (; f, x, y, dy) = deepcopy(scen) f! = f @@ -207,7 +163,7 @@ end ## Gradient -function test_type_gradient_allocating(ba::AbstractADType, scen::Scenario) +function test_type(::GradientAllocating, ba::AbstractADType, scen::Scenario) (; f, x, dx) = deepcopy(scen) grad_in = zero(dx) @@ -226,7 +182,7 @@ end ## Jacobian -function test_type_jacobian_allocating(ba::AbstractADType, scen::Scenario) +function test_type(::JacobianAllocating, ba::AbstractADType, scen::Scenario) (; f, x, y) = deepcopy(scen) jac_in = zeros(eltype(y), length(y), length(x)) @@ -243,7 +199,7 @@ function test_type_jacobian_allocating(ba::AbstractADType, scen::Scenario) @test_opt jacobian(ba, f, x) end -function test_type_jacobian_mutating(ba::AbstractADType, scen::Scenario) +function test_type(::JacobianMutating, ba::AbstractADType, scen::Scenario) Bool(supports_mutation(ba)) || return nothing (; f, x, y) = deepcopy(scen) f! = f @@ -256,7 +212,7 @@ end ## Second derivative -function test_type_second_derivative_allocating(ba::AbstractADType, scen::Scenario) +function test_type(::SecondDerivativeAllocating, ba::AbstractADType, scen::Scenario) (; f, x) = deepcopy(scen) @test_call value_derivative_and_second_derivative(ba, f, x) @@ -268,7 +224,7 @@ end ## Hessian-vector product -function test_type_hessian_vector_product_allocating(ba::AbstractADType, scen::Scenario) +function test_type(::HessianVectorProductAllocating, ba::AbstractADType, scen::Scenario) Bool(supports_hvp(ba)) || return nothing (; f, x, dx) = deepcopy(scen) grad_in = zero(dx) @@ -301,7 +257,7 @@ end ## Hessian -function test_type_hessian_allocating(ba::AbstractADType, scen::Scenario) +function test_type(::HessianAllocating, ba::AbstractADType, scen::Scenario) (; f, x, dx) = deepcopy(scen) grad_in = zero(dx) hess_in = zeros(eltype(x), length(x), length(x)) @@ -323,4 +279,8 @@ function test_type_hessian_allocating(ba::AbstractADType, scen::Scenario) @test_opt ignored_modules = (LinearAlgebra,) hessian(ba, f, x) end +function test_type(op::AbstractOperator, ba::AbstractADType, scen::Scenario) + throw(ArgumentError("Invalid operator to test: $op")) end + +end #module diff --git a/src/DifferentiationTest/DifferentiationTest.jl b/src/DifferentiationTest/DifferentiationTest.jl index 12f73ebd2..72837dd84 100644 --- a/src/DifferentiationTest/DifferentiationTest.jl +++ b/src/DifferentiationTest/DifferentiationTest.jl @@ -15,7 +15,10 @@ using ..DifferentiationInterface: supports_pushforward, supports_pullback, supports_hvp, - zero! + zero!, + MutationBehavior, + MutationSupported, + MutationNotSupported using ADTypes using ADTypes: AbstractADType, @@ -26,6 +29,7 @@ using ADTypes: using DocStringExtensions using Test: @testset, @test +include("operator_traits.jl") include("scenario.jl") include("benchmark.jl") include("call_count.jl") @@ -36,7 +40,7 @@ include("printing.jl") export backend_string export Scenario, default_scenarios -export allocating, mutating, scalar_scalar, scalar_array, array_scalar, array_array +export isallocating, ismutating export BenchmarkData, record! export test_operators export AutoZeroForward, AutoZeroReverse diff --git a/src/DifferentiationTest/call_count.jl b/src/DifferentiationTest/call_count.jl index 6ec71dc61..3b248f73b 100644 --- a/src/DifferentiationTest/call_count.jl +++ b/src/DifferentiationTest/call_count.jl @@ -17,74 +17,14 @@ end function test_call_count( backends::Vector{<:AbstractADType}, - operators::Vector{Symbol}, + operators::Vector{<:AbstractOperator}, scenarios::Vector{<:Scenario}; ) - @testset verbose = true "Correctness" begin + @testset verbose = true "Call count" begin @testset verbose = true "$(backend_string(backend))" for backend in backends @testset "$op" for op in operators - if op == :pushforward_allocating - @testset "$s" for s in allocating(scenarios) - test_call_count_pushforward_allocating(backend, s) - end - elseif op == :pushforward_mutating - @testset "$s" for s in mutating(scenarios) - test_call_count_pushforward_mutating(backend, s) - end - - elseif op == :pullback_allocating - @testset "$s" for s in allocating(scenarios) - test_call_count_pullback_allocating(backend, s) - end - elseif op == :pullback_mutating - @testset "$s" for s in mutating(scenarios) - test_call_count_pullback_mutating(backend, s) - end - - elseif op == :derivative_allocating - @testset "$s" for s in allocating(scalar_scalar(scenarios)) - test_call_count_derivative_allocating(backend, s) - end - - elseif op == :multiderivative_allocating - @testset "$s" for s in allocating(scalar_array(scenarios)) - test_call_count_multiderivative_allocating(backend, s) - end - elseif op == :multiderivative_mutating - @testset "$s" for s in mutating(scalar_array(scenarios)) - test_call_count_multiderivative_mutating(backend, s) - end - - elseif op == :gradient_allocating - @testset "$s" for s in allocating(array_scalar(scenarios)) - test_call_count_gradient_allocating(backend, s) - end - - elseif op == :jacobian_allocating - @testset "$s" for s in allocating(array_array(scenarios)) - test_call_count_jacobian_allocating(backend, s) - end - elseif op == :jacobian_mutating - @testset "$s" for s in mutating(array_array(scenarios)) - test_call_count_jacobian_mutating(backend, s) - end - - elseif op == :second_derivative_allocating - @testset "$s" for s in allocating(scalar_scalar(scenarios)) - test_call_count_second_derivative_allocating(backend, s) - end - - elseif op == :hessian_vector_product_allocating - @testset "$s" for s in allocating(array_scalar(scenarios)) - test_call_count_hessian_vector_product_allocating(backend, s) - end - elseif op == :hessian_allocating - @testset "$s" for s in allocating(array_scalar(scenarios)) - test_call_count_hessian_allocating(backend, s) - end - - else - throw(ArgumentError("Invalid operator to test: `:$op`")) + @testset "$s" for s in compatible_scenarios(op, scenarios) + test_call_count(op, backend, s) end end end @@ -93,7 +33,7 @@ end ## Pushforward -function test_call_count_pushforward_allocating(ba::AbstractADType, scen::Scenario) +function test_call_count(::PushforwardAllocating, ba::AbstractADType, scen::Scenario) Bool(supports_pushforward(ba)) || return nothing (; f, x, dx) = deepcopy(scen) cc = CallCounter(f) @@ -103,7 +43,7 @@ function test_call_count_pushforward_allocating(ba::AbstractADType, scen::Scenar end end -function test_call_count_pushforward_mutating(ba::AbstractADType, scen::Scenario) +function test_call_count(::PushforwardMutating, ba::AbstractADType, scen::Scenario) Bool(supports_pushforward(ba)) || return nothing Bool(supports_mutation(ba)) || return nothing (; f, x, y, dx, dy) = deepcopy(scen) @@ -118,7 +58,7 @@ end ## Pullback -function test_call_count_pullback_allocating(ba::AbstractADType, scen::Scenario) +function test_call_count(::PullbackAllocating, ba::AbstractADType, scen::Scenario) Bool(supports_pullback(ba)) || return nothing (; f, x, y, dy) = deepcopy(scen) cc = CallCounter(f) @@ -128,7 +68,7 @@ function test_call_count_pullback_allocating(ba::AbstractADType, scen::Scenario) end end -function test_call_count_pullback_mutating(ba::AbstractADType, scen::Scenario) +function test_call_count(::PullbackMutating, ba::AbstractADType, scen::Scenario) Bool(supports_pullback(ba)) || return nothing Bool(supports_mutation(ba)) || return nothing (; f, x, y, dx, dy) = deepcopy(scen) @@ -143,8 +83,8 @@ end ## Derivative -function test_call_count_derivative_allocating( - ba::AbstractADType, scen::Scenario{<:Any,<:Number} +function test_call_count( + ::DerivativeAllocating, ba::AbstractADType, scen::Scenario{<:Any,<:Number} ) (; f, x, y) = deepcopy(scen) cc = CallCounter(f) @@ -154,8 +94,10 @@ end ## Multiderivative -function test_call_count_multiderivative_allocating( - ba::AbstractADType, scen::Scenario{<:Any,<:Number,<:AbstractArray} +function test_call_count( + ::MultiderivativeAllocating, + ba::AbstractADType, + scen::Scenario{<:Any,<:Number,<:AbstractArray}, ) (; f, x, y) = deepcopy(scen) cc1 = CallCounter(f) @@ -171,8 +113,10 @@ function test_call_count_multiderivative_allocating( end end -function test_call_count_multiderivative_mutating( - ba::AbstractADType, scen::Scenario{<:Any,<:Number,<:AbstractArray} +function test_call_count( + ::MultiderivativeMutating, + ba::AbstractADType, + scen::Scenario{<:Any,<:Number,<:AbstractArray}, ) Bool(supports_mutation(ba)) || return nothing (; f, x, y, dy) = deepcopy(scen) @@ -189,8 +133,8 @@ end ## Gradient -function test_call_count_gradient_allocating( - ba::AbstractADType, scen::Scenario{<:Any,<:AbstractArray,<:Number} +function test_call_count( + ::GradientAllocating, ba::AbstractADType, scen::Scenario{<:Any,<:AbstractArray,<:Number} ) (; f, x, y) = deepcopy(scen) cc1 = CallCounter(f) @@ -208,8 +152,10 @@ end ## Jacobian -function test_call_count_jacobian_allocating( - ba::AbstractADType, scen::Scenario{<:Any,<:AbstractArray,<:AbstractArray} +function test_call_count( + ::JacobianAllocating, + ba::AbstractADType, + scen::Scenario{<:Any,<:AbstractArray,<:AbstractArray}, ) (; f, x, y) = deepcopy(scen) cc1 = CallCounter(f) @@ -225,8 +171,10 @@ function test_call_count_jacobian_allocating( end end -function test_call_count_jacobian_mutating( - ba::AbstractADType, scen::Scenario{<:Any,<:AbstractArray,<:AbstractArray} +function test_call_count( + ::JacobianMutating, + ba::AbstractADType, + scen::Scenario{<:Any,<:AbstractArray,<:AbstractArray}, ) Bool(supports_mutation(ba)) || return nothing (; f, x, y) = deepcopy(scen) @@ -243,8 +191,10 @@ end ## Second derivative -function test_call_count_second_derivative_allocating( - ba::AbstractADType, scen::Scenario{<:Any,<:Number,<:Number} +function test_call_count( + ::SecondDerivativeAllocating, + ba::AbstractADType, + scen::Scenario{<:Any,<:Number,<:Number}, ) (; f, x, y) = deepcopy(scen) cc = CallCounter(f) @@ -254,8 +204,10 @@ end ## Hessian-vector product -function test_call_count_hessian_vector_product_allocating( - ba::AbstractADType, scen::Scenario{<:Any,<:AbstractArray,<:Number} +function test_call_count( + ::HessianVectorProductAllocating, + ba::AbstractADType, + scen::Scenario{<:Any,<:AbstractArray,<:Number}, ) Bool(supports_hvp(ba)) || return nothing (; f, x, dx) = deepcopy(scen) @@ -269,8 +221,8 @@ end ## Hessian -function test_call_count_hessian_allocating( - ba::AbstractADType, scen::Scenario{<:Any,<:AbstractArray,<:Number} +function test_call_count( + ::HessianAllocating, ba::AbstractADType, scen::Scenario{<:Any,<:AbstractArray,<:Number} ) (; f, x, y) = deepcopy(scen) cc1 = CallCounter(f) @@ -280,3 +232,7 @@ function test_call_count_hessian_allocating( @test cc1.count[] <= 2 + length(x) @test cc2.count[] <= 2 + length(x) end + +function test_call_count(op::AbstractOperator, ba::AbstractADType, scen::Scenario) + throw(ArgumentError("Invalid operator to test: $op")) +end diff --git a/src/DifferentiationTest/default_scenarios.jl b/src/DifferentiationTest/default_scenarios.jl index 79e031ee6..4e3f7cb0c 100644 --- a/src/DifferentiationTest/default_scenarios.jl +++ b/src/DifferentiationTest/default_scenarios.jl @@ -1,52 +1,52 @@ const SCALING_VEC = Vector(1:12) const SCALING_MAT = Matrix((1:3) .* transpose(1:4)) -f_scalar_scalar(x::Number)::Number = sin(x) +f_scalar_to_scalar(x::Number)::Number = sin(x) -function f_scalar_vector(x::Number)::AbstractVector +function f_scalar_to_vector(x::Number)::AbstractVector return sin.(SCALING_VEC .* x) # output size 12 end -function f!_scalar_vector(y::AbstractVector, x::Number) +function f!_scalar_to_vector(y::AbstractVector, x::Number) for i in eachindex(y) y[i] = sin(i * x) end return nothing end -function f_scalar_matrix(x::Number)::AbstractMatrix +function f_scalar_to_matrix(x::Number)::AbstractMatrix return sin.(SCALING_MAT .* x) # output size (3, 4) end -function f!_scalar_matrix(y::AbstractMatrix, x::Number) +function f!_scalar_to_matrix(y::AbstractMatrix, x::Number) for i in axes(y, 1), j in axes(y, 2) y[i, j] = sin(i * j * x) end return nothing end -f_vector_scalar(x::AbstractVector)::Number = sum(sin, x) -f_matrix_scalar(x::AbstractMatrix)::Number = sum(sin, x) +f_vector_to_scalar(x::AbstractVector)::Number = sum(sin, x) +f_matrix_to_scalar(x::AbstractMatrix)::Number = sum(sin, x) -f_vector_vector(x::AbstractVector)::AbstractVector = vcat(sin.(x), cos.(x)) +f_vector_to_vector(x::AbstractVector)::AbstractVector = vcat(sin.(x), cos.(x)) -function f!_vector_vector(y::AbstractVector, x::AbstractVector) +function f!_vector_to_vector(y::AbstractVector, x::AbstractVector) y[1:length(x)] .= sin.(x) y[(length(x) + 1):(2length(x))] .= cos.(x) return nothing end -f_vector_matrix(x::AbstractVector)::AbstractMatrix = hcat(sin.(x), cos.(x)) +f_vector_to_matrix(x::AbstractVector)::AbstractMatrix = hcat(sin.(x), cos.(x)) -function f!_vector_matrix(y::AbstractMatrix, x::AbstractVector) +function f!_vector_to_matrix(y::AbstractMatrix, x::AbstractVector) y[:, 1] .= sin.(x) y[:, 2] .= cos.(x) return nothing end -f_matrix_vector(x::AbstractMatrix)::AbstractVector = vcat(vec(sin.(x)), vec(cos.(x))) +f_matrix_to_vector(x::AbstractMatrix)::AbstractVector = vcat(vec(sin.(x)), vec(cos.(x))) -function f!_matrix_vector(y::AbstractVector, x::AbstractMatrix) +function f!_matrix_to_vector(y::AbstractVector, x::AbstractMatrix) for i in eachindex(IndexLinear(), x) y[i] = sin(x[i]) y[length(x) + i] = cos(x[i]) @@ -54,9 +54,9 @@ function f!_matrix_vector(y::AbstractVector, x::AbstractMatrix) return nothing end -f_matrix_matrix(x::AbstractMatrix)::AbstractMatrix = hcat(vec(sin.(x)), vec(cos.(x))) +f_matrix_to_matrix(x::AbstractMatrix)::AbstractMatrix = hcat(vec(sin.(x)), vec(cos.(x))) -function f!_matrix_matrix(y::AbstractMatrix, x::AbstractMatrix) +function f!_matrix_to_matrix(y::AbstractMatrix, x::AbstractMatrix) for i in eachindex(IndexLinear(), x) y[i, 1] = sin(x[i]) y[i, 2] = cos(x[i]) @@ -66,27 +66,27 @@ end function default_scenarios_allocating() scenarios = [ - Scenario(f_scalar_scalar, 2.0), - Scenario(f_scalar_vector, 2.0), - Scenario(f_scalar_matrix, 2.0), - Scenario(f_vector_scalar, Vector{Float64}(1:12)), - Scenario(f_matrix_scalar, Matrix{Float64}(reshape(1:12, 3, 4))), - Scenario(f_vector_vector, Vector{Float64}(1:12)), - Scenario(f_vector_matrix, Vector{Float64}(1:12)), - Scenario(f_matrix_vector, Matrix{Float64}(reshape(1:12, 3, 4))), - Scenario(f_matrix_matrix, Matrix{Float64}(reshape(1:12, 3, 4))), + Scenario(f_scalar_to_scalar, 2.0), + Scenario(f_scalar_to_vector, 2.0), + Scenario(f_scalar_to_matrix, 2.0), + Scenario(f_vector_to_scalar, Vector{Float64}(1:12)), + Scenario(f_matrix_to_scalar, Matrix{Float64}(reshape(1:12, 3, 4))), + Scenario(f_vector_to_vector, Vector{Float64}(1:12)), + Scenario(f_vector_to_matrix, Vector{Float64}(1:12)), + Scenario(f_matrix_to_vector, Matrix{Float64}(reshape(1:12, 3, 4))), + Scenario(f_matrix_to_matrix, Matrix{Float64}(reshape(1:12, 3, 4))), ] return scenarios end function default_scenarios_mutating() scenarios = [ - Scenario(f!_scalar_vector, 2.0, (12,)), - Scenario(f!_scalar_matrix, 2.0, (3, 4)), - Scenario(f!_vector_vector, Vector{Float64}(1:12), (24,)), - Scenario(f!_vector_matrix, Vector{Float64}(1:12), (12, 2)), - Scenario(f!_matrix_vector, Matrix{Float64}(reshape(1:12, 3, 4)), (24,)), - Scenario(f!_matrix_matrix, Matrix{Float64}(reshape(1:12, 3, 4)), (12, 2)), + Scenario(f!_scalar_to_vector, 2.0, (12,)), + Scenario(f!_scalar_to_matrix, 2.0, (3, 4)), + Scenario(f!_vector_to_vector, Vector{Float64}(1:12), (24,)), + Scenario(f!_vector_to_matrix, Vector{Float64}(1:12), (12, 2)), + Scenario(f!_matrix_to_vector, Matrix{Float64}(reshape(1:12, 3, 4)), (24,)), + Scenario(f!_matrix_to_matrix, Matrix{Float64}(reshape(1:12, 3, 4)), (12, 2)), ] return scenarios end diff --git a/src/DifferentiationTest/operator_traits.jl b/src/DifferentiationTest/operator_traits.jl new file mode 100644 index 000000000..ed761bbeb --- /dev/null +++ b/src/DifferentiationTest/operator_traits.jl @@ -0,0 +1,115 @@ +abstract type AbstractOperator end +abstract type AbstractFirstOrderOperator <: AbstractOperator end +abstract type AbstractSecondOrderOperator <: AbstractOperator end + +# First-Order Operators +struct Pullback{M<:MutationBehavior} <: AbstractFirstOrderOperator end +struct Pushforward{M<:MutationBehavior} <: AbstractFirstOrderOperator end +struct Gradient{M<:MutationBehavior} <: AbstractFirstOrderOperator end +struct Multiderivative{M<:MutationBehavior} <: AbstractFirstOrderOperator end +struct Jacobian{M<:MutationBehavior} <: AbstractFirstOrderOperator end +struct Derivative{M<:MutationBehavior} <: AbstractFirstOrderOperator end + +const PullbackAllocating = Pullback{MutationNotSupported} +const PullbackMutating = Pullback{MutationSupported} +const PushforwardAllocating = Pushforward{MutationNotSupported} +const PushforwardMutating = Pushforward{MutationSupported} +const GradientAllocating = Gradient{MutationNotSupported} +const GradientMutating = Gradient{MutationSupported} +const MultiderivativeAllocating = Multiderivative{MutationNotSupported} +const MultiderivativeMutating = Multiderivative{MutationSupported} +const JacobianAllocating = Jacobian{MutationNotSupported} +const JacobianMutating = Jacobian{MutationSupported} +const DerivativeAllocating = Derivative{MutationNotSupported} +const DerivativeMutating = Derivative{MutationSupported} + +# Second-order operators +struct SecondDerivative{M<:MutationBehavior} <: AbstractSecondOrderOperator end +struct Hessian{M<:MutationBehavior} <: AbstractSecondOrderOperator end +struct HessianVectorProduct{M<:MutationBehavior} <: AbstractSecondOrderOperator end + +const SecondDerivativeAllocating = SecondDerivative{MutationNotSupported} +const SecondDerivativeMutating = SecondDerivative{MutationSupported} +const HessianAllocating = Hessian{MutationNotSupported} +const HessianMutating = Hessian{MutationSupported} +const HessianVectorProductAllocating = HessianVectorProduct{MutationNotSupported} +const HessianVectorProductMutating = HessianVectorProduct{MutationSupported} + +## Utilities +# order +isfirstorder(::AbstractOperator) = false +isfirstorder(::AbstractFirstOrderOperator) = true + +issecondorder(::AbstractOperator) = false +issecondorder(::AbstractSecondOrderOperator) = true + +# allocations +ismutating(::Type{<:MutationBehavior}) = false +ismutating(::Type{MutationSupported}) = true + +ismutating(::Pullback{M}) where {M} = ismutating(M) +ismutating(::Pushforward{M}) where {M} = ismutating(M) +ismutating(::Gradient{M}) where {M} = ismutating(M) +ismutating(::Multiderivative{M}) where {M} = ismutating(M) +ismutating(::Derivative{M}) where {M} = ismutating(M) +ismutating(::Jacobian{M}) where {M} = ismutating(M) +ismutating(::SecondDerivative{M}) where {M} = ismutating(M) +ismutating(::Hessian{M}) where {M} = ismutating(M) +ismutating(::HessianVectorProduct{M}) where {M} = ismutating(M) + +isallocating(op) = !ismutating(op) + +# input-output compatibility +iscompatible(op::AbstractOperator, x, y) = false +iscompatible(op::Pullback, x, y) = true +iscompatible(op::Pushforward, x, y) = true + +iscompatible(op::Gradient, x::AbstractArray, y::Number) = true +iscompatible(op::Multiderivative, x::Number, y::AbstractArray) = true +iscompatible(op::Derivative, x::Number, y::Number) = true +iscompatible(op::Jacobian, x::AbstractArray, y::AbstractArray) = true +iscompatible(op::SecondDerivative, x::Number, y::Number) = true +iscompatible(op::Hessian, x::AbstractArray, y::Number) = true +iscompatible(op::HessianVectorProduct, x::AbstractArray, y::Number) = true + +## Pretty-printing +alloc_string(T::Type{<:MutationBehavior}) = ismutating(T) ? "mutating" : "allocating" + +Base.string(::Pullback{M}) where {M} = "Pullback, $(alloc_string(M))" +Base.string(::Pushforward{M}) where {M} = "Pushforward, $(alloc_string(M))" +Base.string(::Gradient{M}) where {M} = "Gradient, $(alloc_string(M))" +Base.string(::Multiderivative{M}) where {M} = "Multiderivative, $(alloc_string(M))" +Base.string(::Derivative{M}) where {M} = "Derivative, $(alloc_string(M))" +Base.string(::Jacobian{M}) where {M} = "Jacobian, $(alloc_string(M))" +Base.string(::SecondDerivative{M}) where {M} = "Second derivative, $(alloc_string(M))" +Base.string(::Hessian{M}) where {M} = "Hessian, $(alloc_string(M))" +Base.string(::HessianVectorProduct{M}) where {M} = "Hessian-vector product, $(alloc_string(M))" + +## Convert symbols to traits +const OPERATOR_SYMBOL_TO_TRAIT = Dict( + :pushforward_allocating => PushforwardAllocating(), + :pushforward_mutating => PushforwardMutating(), + :pullback_allocating => PullbackAllocating(), + :pullback_mutating => PullbackMutating(), + :multiderivative_allocating => MultiderivativeAllocating(), + :multiderivating_mutating => MultiderivativeMutating(), + :derivative_allocating => DerivativeAllocating(), + :derivative_mutating => DerivativeMutating(), + :gradient_allocating => GradientAllocating(), + :gradient_mutating => GradientMutating(), + :jacobian_allocating => JacobianAllocating(), + :jacobian_mutating => JacobianMutating(), + :second_derivative_allocating => SecondDerivativeAllocating(), + :second_derivative_mutating => SecondDerivativeMutating(), + :hessian_allocating => HessianAllocating(), + :hessian_mutating => HessianMutating(), + :hessian_vector_product_allocating => HessianVectorProductAllocating(), + :hessian_vector_product_mutating => HessianVectorProductMutating(), +) + +operator_trait(op::AbstractOperator) = op +function operator_trait(sym::Symbol) + !haskey(OPERATOR_SYMBOL_TO_TRAIT, sym) && + throw(ArgumentError("Invalid operator symbol: $sym")) + return OPERATOR_SYMBOL_TO_TRAIT[sym] +end diff --git a/src/DifferentiationTest/scenario.jl b/src/DifferentiationTest/scenario.jl index a0375fb98..1d0e773e1 100644 --- a/src/DifferentiationTest/scenario.jl +++ b/src/DifferentiationTest/scenario.jl @@ -25,6 +25,19 @@ end Base.string(scen::Scenario) = "$(string(scen.f)): $(typeof(scen.x)) -> $(typeof(scen.y))" +ismutating(s::Scenario) = s.mutating +isallocating(s::Scenario) = !ismutating(s) + +## Check operator compatibility +function iscompatible(op::AbstractOperator, s::Scenario) + return ismutating(op) == ismutating(s) && iscompatible(op, s.x, s.y) +end + +function compatible_scenarios(op::AbstractOperator, scs::AbstractVector{Scenario}) + return filter(s -> iscompatible(op, s), scs) +end + +## Scenario constructors similar_random(z::Number) = randn(eltype(z)) function similar_random(z::AbstractArray) @@ -47,30 +60,3 @@ function Scenario(f!, x::Union{Number,AbstractArray}, s::NTuple{N,<:Integer}) wh dy = similar_random(y) return Scenario(; f=f!, x, y, dx, dy, mutating=true) end - -allocating(scenarios::Vector{<:Scenario}) = filter(s -> !s.mutating, scenarios) -mutating(scenarios::Vector{<:Scenario}) = filter(s -> s.mutating, scenarios) - -function scalar_scalar(scenarios::Vector{<:Scenario}) - return filter(scenarios) do s - typeof(s.x) <: Number && typeof(s.y) <: Number - end -end - -function scalar_array(scenarios::Vector{<:Scenario}) - return filter(scenarios) do s - typeof(s.x) <: Number && typeof(s.y) <: AbstractArray - end -end - -function array_scalar(scenarios::Vector{<:Scenario}) - return filter(scenarios) do s - typeof(s.x) <: AbstractArray && typeof(s.y) <: Number - end -end - -function array_array(scenarios::Vector{<:Scenario}) - return filter(scenarios) do s - typeof(s.x) <: AbstractArray && typeof(s.y) <: AbstractArray - end -end diff --git a/src/DifferentiationTest/test_operators.jl b/src/DifferentiationTest/test_operators.jl index 7aa90c97e..8815dffc0 100644 --- a/src/DifferentiationTest/test_operators.jl +++ b/src/DifferentiationTest/test_operators.jl @@ -2,54 +2,49 @@ test_correctness(args...; kwargs...) = error("Please load ForwardDiff.jl") test_type_stability(args...; kwargs...) = error("Please load JET.jl") const FIRST_ORDER_OPERATORS = [ - :pushforward_allocating, - :pushforward_mutating, - :pullback_allocating, - :pullback_mutating, - :derivative_allocating, - :multiderivative_allocating, - :multiderivative_mutating, - :gradient_allocating, - :jacobian_allocating, - :jacobian_mutating, + PushforwardAllocating(), + PushforwardMutating(), + PullbackAllocating(), + PullbackMutating(), + MultiderivativeAllocating(), + MultiderivativeMutating(), + DerivativeAllocating(), + # DerivativeMutating(), + GradientAllocating(), + # GradientMutating(), + JacobianAllocating(), + JacobianMutating(), ] const SECOND_ORDER_OPERATORS = [ - :second_derivative_allocating, :hessian_vector_product_allocating, :hessian_allocating + SecondDerivativeAllocating(), + # SecondDerivativeMutating(), + HessianAllocating(), + # HessianMutating(), + HessianVectorProductAllocating(), + # HessianVectorProductMutating(), ] +const ALL_OPERATORS = vcat(FIRST_ORDER_OPERATORS, SECOND_ORDER_OPERATORS) + function filter_operators( - operators::Vector{Symbol}; + operators::Vector{<:AbstractOperator}; first_order::Bool, second_order::Bool, allocating::Bool, mutating::Bool, - excluded::Vector{Symbol}, + excluded::Vector{<:AbstractOperator}, ) - if !first_order - operators = setdiff(operators, FIRST_ORDER_OPERATORS) - end - if !second_order - operators = setdiff(operators, SECOND_ORDER_OPERATORS) - end - if !allocating - operators = filter(op -> !endswith(string(op), "allocating"), operators) - end - if !mutating - operators = filter(op -> !endswith(string(op), "mutating"), operators) - end + !first_order && (operators = filter(!isfirstorder, operators)) + !second_order && (operators = filter(!issecondorder, operators)) + !allocating && (operators = filter(!isallocating, operators)) + !mutating && (operators = filter(!ismutating, operators)) operators = filter(op -> !in(op, excluded), operators) return operators end """ - test_operators( - backends, [operators, scenarios]; - correctness, type_stability, benchmark, allocations, - input_type, output_type, - first_order, second_order, allocating, mutating, - excluded, - ) + test_operators(backends, [operators, scenarios]; [kwargs...]) Cross-test a list of `backends` for a list of `operators` on a list of `scenarios.` @@ -57,7 +52,7 @@ Return `nothing`, except when `benchmark=true`. # Default arguments -- `operators`: defaults to all of them +- `operators`: defaults to all operators - `scenarios`: defaults to a set of default scenarios # Keyword arguments @@ -77,7 +72,7 @@ Return `nothing`, except when `benchmark=true`. """ function test_operators( backends::Vector{<:AbstractADType}, - operators::Vector{Symbol}=vcat(FIRST_ORDER_OPERATORS, SECOND_ORDER_OPERATORS), + operators::Vector{<:AbstractOperator}=ALL_OPERATORS, scenarios::Vector{<:Scenario}=default_scenarios(); correctness::Bool=true, type_stability::Bool=true, @@ -90,8 +85,9 @@ function test_operators( second_order=true, allocating=true, mutating=true, - excluded::Vector{Symbol}=Symbol[], + excluded::Union{Vector{<:AbstractOperator},Vector{Symbol}}=AbstractOperator[], ) + excluded = operator_trait.(excluded) scenarios = filter(scenarios) do scen typeof(scen.x) <: input_type && typeof(scen.y) <: output_type end @@ -130,3 +126,10 @@ Shortcut for a single backend. function test_operators(backend::AbstractADType, args...; kwargs...) return test_operators([backend], args...; kwargs...) end + +function test_operators( + backend::AbstractADType, operators::Vector{Symbol}, args...; kwargs... +) + operators = operator_trait.(operators) + return test_operators([backend], operators, args...; kwargs...) +end diff --git a/test/runtests.jl b/test/runtests.jl index e2f21ab3c..af8de2b4d 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -22,14 +22,17 @@ using Zygote: Zygote @testset verbose = true "DifferentiationInterface.jl" begin @testset verbose = true "Formal tests" begin @testset "Aqua" begin + @info "Running Aqua.jl tests..." Aqua.test_all(DifferentiationInterface; ambiguities=false) end @testset "JuliaFormatter" begin + @info "Running JuliaFormatter test..." @test JuliaFormatter.format( DifferentiationInterface; verbose=false, overwrite=false ) end @testset "JET" begin + @info "Running JET tests..." JET.test_package(DifferentiationInterface; target_defined_modules=true) end end @@ -45,41 +48,53 @@ using Zygote: Zygote @testset verbose = true "First order" begin @testset "ChainRules (reverse)" begin + @info "Running ChainRules (reverse) tests..." include("chainrules_reverse.jl") end @testset "Diffractor (forward)" begin + @info "Running Diffractor (forward) tests..." include("diffractor.jl") end @testset "Enzyme (forward)" begin + @info "Running Enzyme (forward) tests..." include("enzyme_forward.jl") end @testset "Enzyme (reverse)" begin + @info "Running Enzyme (reverse) tests..." include("enzyme_reverse.jl") end @testset "FastDifferentiation" begin + @info "Running FastDifferentiation tests..." include("fastdifferentiation.jl") end @testset "FiniteDiff" begin + @info "Running FiniteDiff tests..." include("finitediff.jl") end @testset "ForwardDiff" begin + @info "Running ForwardDiff tests..." include("forwarddiff.jl") end @testset "PolyesterForwardDiff" begin + @info "Running PolyesterForwardDiff tests..." include("polyesterforwarddiff.jl") end @testset "ReverseDiff" begin + @info "Running ReverseDiff tests..." include("reversediff.jl") end @testset "Tracker" begin + @info "Running Tracker tests..." include("tracker.jl") end @testset "Zygote" begin + @info "Running Zygote tests..." include("zygote.jl") end end @testset verbose = true "Second order" begin + @info "Running second order tests..." include("second_order.jl") end end