diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/onearg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/onearg.jl index f3ebc6b1c..a3b56e1aa 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/onearg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/onearg.jl @@ -305,24 +305,30 @@ function DI.value_and_gradient!( f::F, grad, prep::ForwardDiffGradientPrep, - ::AutoForwardDiff, + backend::AutoForwardDiff, x, contexts::Vararg{Constant,C}, ) where {F,C} fc = with_contexts(f, contexts...) result = DiffResult(zero(eltype(x)), (grad,)) - result = gradient!(result, fc, x, prep.config) + CHK = tag_type(backend) === Nothing + result = gradient!(result, fc, x, prep.config, Val(CHK)) y = DR.value(result) grad === DR.gradient(result) || copyto!(grad, DR.gradient(result)) return y, grad end function DI.value_and_gradient( - f::F, prep::ForwardDiffGradientPrep, ::AutoForwardDiff, x, contexts::Vararg{Constant,C} + f::F, + prep::ForwardDiffGradientPrep, + backend::AutoForwardDiff, + x, + contexts::Vararg{Constant,C}, ) where {F,C} fc = with_contexts(f, contexts...) result = GradientResult(x) - result = gradient!(result, fc, x, prep.config) + CHK = tag_type(backend) === Nothing + result = gradient!(result, fc, x, prep.config, Val(CHK)) return DR.value(result), DR.gradient(result) end @@ -330,19 +336,25 @@ function DI.gradient!( f::F, grad, prep::ForwardDiffGradientPrep, - ::AutoForwardDiff, + backend::AutoForwardDiff, x, contexts::Vararg{Constant,C}, ) where {F,C} fc = with_contexts(f, contexts...) - return gradient!(grad, fc, x, prep.config) + CHK = tag_type(backend) === Nothing + return gradient!(grad, fc, x, prep.config, Val(CHK)) end function DI.gradient( - f::F, prep::ForwardDiffGradientPrep, ::AutoForwardDiff, x, contexts::Vararg{Constant,C} + f::F, + prep::ForwardDiffGradientPrep, + backend::AutoForwardDiff, + x, + contexts::Vararg{Constant,C}, ) where {F,C} fc = with_contexts(f, contexts...) - return gradient(fc, x, prep.config) + CHK = tag_type(backend) === Nothing + return gradient(fc, x, prep.config, Val(CHK)) end ## Jacobian @@ -422,14 +434,15 @@ function DI.value_and_jacobian!( f::F, jac, prep::ForwardDiffOneArgJacobianPrep, - ::AutoForwardDiff, + backend::AutoForwardDiff, x, contexts::Vararg{Constant,C}, ) where {F,C} fc = with_contexts(f, contexts...) y = fc(x) result = DiffResult(y, (jac,)) - result = jacobian!(result, fc, x, prep.config) + CHK = tag_type(backend) === Nothing + result = jacobian!(result, fc, x, prep.config, Val(CHK)) y = DR.value(result) jac === DR.jacobian(result) || copyto!(jac, DR.jacobian(result)) return y, jac @@ -438,35 +451,38 @@ end function DI.value_and_jacobian( f::F, prep::ForwardDiffOneArgJacobianPrep, - ::AutoForwardDiff, + backend::AutoForwardDiff, x, contexts::Vararg{Constant,C}, ) where {F,C} fc = with_contexts(f, contexts...) - return fc(x), jacobian(fc, x, prep.config) + CHK = tag_type(backend) === Nothing + return fc(x), jacobian(fc, x, prep.config, Val(CHK)) end function DI.jacobian!( f::F, jac, prep::ForwardDiffOneArgJacobianPrep, - ::AutoForwardDiff, + backend::AutoForwardDiff, x, contexts::Vararg{Constant,C}, ) where {F,C} fc = with_contexts(f, contexts...) - return jacobian!(jac, fc, x, prep.config) + CHK = tag_type(backend) === Nothing + return jacobian!(jac, fc, x, prep.config, Val(CHK)) end function DI.jacobian( f::F, prep::ForwardDiffOneArgJacobianPrep, - ::AutoForwardDiff, + backend::AutoForwardDiff, x, contexts::Vararg{Constant,C}, ) where {F,C} fc = with_contexts(f, contexts...) - return jacobian(fc, x, prep.config) + CHK = tag_type(backend) === Nothing + return jacobian(fc, x, prep.config, Val(CHK)) end ## Second derivative @@ -681,19 +697,25 @@ function DI.hessian!( f::F, hess, prep::ForwardDiffHessianPrep, - ::AutoForwardDiff, + backend::AutoForwardDiff, x, contexts::Vararg{Constant,C}, ) where {F,C} fc = with_contexts(f, contexts...) - return hessian!(hess, fc, x, prep.array_config) + CHK = tag_type(backend) === Nothing + return hessian!(hess, fc, x, prep.array_config, Val(CHK)) end function DI.hessian( - f::F, prep::ForwardDiffHessianPrep, ::AutoForwardDiff, x, contexts::Vararg{Constant,C} + f::F, + prep::ForwardDiffHessianPrep, + backend::AutoForwardDiff, + x, + contexts::Vararg{Constant,C}, ) where {F,C} fc = with_contexts(f, contexts...) - return hessian(fc, x, prep.array_config) + CHK = tag_type(backend) === Nothing + return hessian(fc, x, prep.array_config, Val(CHK)) end function DI.value_gradient_and_hessian!( @@ -701,13 +723,14 @@ function DI.value_gradient_and_hessian!( grad, hess, prep::ForwardDiffHessianPrep, - ::AutoForwardDiff, + backend::AutoForwardDiff, x, contexts::Vararg{Constant,C}, ) where {F,C} fc = with_contexts(f, contexts...) result = DiffResult(one(eltype(x)), (grad, hess)) - result = hessian!(result, fc, x, prep.result_config) + CHK = tag_type(backend) === Nothing + result = hessian!(result, fc, x, prep.result_config, Val(CHK)) y = DR.value(result) grad === DR.gradient(result) || copyto!(grad, DR.gradient(result)) hess === DR.hessian(result) || copyto!(hess, DR.hessian(result)) @@ -715,10 +738,15 @@ function DI.value_gradient_and_hessian!( end function DI.value_gradient_and_hessian( - f::F, prep::ForwardDiffHessianPrep, ::AutoForwardDiff, x, contexts::Vararg{Constant,C} + f::F, + prep::ForwardDiffHessianPrep, + backend::AutoForwardDiff, + x, + contexts::Vararg{Constant,C}, ) where {F,C} fc = with_contexts(f, contexts...) result = HessianResult(x) - result = hessian!(result, fc, x, prep.result_config) + CHK = tag_type(backend) === Nothing + result = hessian!(result, fc, x, prep.result_config, Val(CHK)) return (DR.value(result), DR.gradient(result), DR.hessian(result)) end diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/twoarg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/twoarg.jl index ce74dd283..f40df3743 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/twoarg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/twoarg.jl @@ -181,13 +181,14 @@ function DI.value_and_derivative( f!::F, y, prep::ForwardDiffTwoArgDerivativePrep, - ::AutoForwardDiff, + backend::AutoForwardDiff, x, contexts::Vararg{Constant,C}, ) where {F,C} fc! = with_contexts(f!, contexts...) result = MutableDiffResult(y, (similar(y),)) - result = derivative!(result, fc!, y, x, prep.config) + CHK = tag_type(backend) === Nothing + result = derivative!(result, fc!, y, x, prep.config, Val(CHK)) return DiffResults.value(result), DiffResults.derivative(result) end @@ -196,13 +197,14 @@ function DI.value_and_derivative!( y, der, prep::ForwardDiffTwoArgDerivativePrep, - ::AutoForwardDiff, + backend::AutoForwardDiff, x, contexts::Vararg{Constant,C}, ) where {F,C} fc! = with_contexts(f!, contexts...) result = MutableDiffResult(y, (der,)) - result = derivative!(result, fc!, y, x, prep.config) + CHK = tag_type(backend) === Nothing + result = derivative!(result, fc!, y, x, prep.config, Val(CHK)) return DiffResults.value(result), DiffResults.derivative(result) end @@ -210,12 +212,13 @@ function DI.derivative( f!::F, y, prep::ForwardDiffTwoArgDerivativePrep, - ::AutoForwardDiff, + backend::AutoForwardDiff, x, contexts::Vararg{Constant,C}, ) where {F,C} fc! = with_contexts(f!, contexts...) - return derivative(fc!, y, x, prep.config) + CHK = tag_type(backend) === Nothing + return derivative(fc!, y, x, prep.config, Val(CHK)) end function DI.derivative!( @@ -223,12 +226,13 @@ function DI.derivative!( y, der, prep::ForwardDiffTwoArgDerivativePrep, - ::AutoForwardDiff, + backend::AutoForwardDiff, x, contexts::Vararg{Constant,C}, ) where {F,C} fc! = with_contexts(f!, contexts...) - return derivative!(der, fc!, y, x, prep.config) + CHK = tag_type(backend) === Nothing + return derivative!(der, fc!, y, x, prep.config, Val(CHK)) end ## Jacobian @@ -308,14 +312,15 @@ function DI.value_and_jacobian( f!::F, y, prep::ForwardDiffTwoArgJacobianPrep, - ::AutoForwardDiff, + backend::AutoForwardDiff, x, contexts::Vararg{Constant,C}, ) where {F,C} fc! = with_contexts(f!, contexts...) jac = similar(y, length(y), length(x)) result = MutableDiffResult(y, (jac,)) - result = jacobian!(result, fc!, y, x, prep.config) + CHK = tag_type(backend) === Nothing + result = jacobian!(result, fc!, y, x, prep.config, Val(CHK)) return DiffResults.value(result), DiffResults.jacobian(result) end @@ -324,13 +329,14 @@ function DI.value_and_jacobian!( y, jac, prep::ForwardDiffTwoArgJacobianPrep, - ::AutoForwardDiff, + backend::AutoForwardDiff, x, contexts::Vararg{Constant,C}, ) where {F,C} fc! = with_contexts(f!, contexts...) result = MutableDiffResult(y, (jac,)) - result = jacobian!(result, fc!, y, x, prep.config) + CHK = tag_type(backend) === Nothing + result = jacobian!(result, fc!, y, x, prep.config, Val(CHK)) return DiffResults.value(result), DiffResults.jacobian(result) end @@ -338,12 +344,13 @@ function DI.jacobian( f!::F, y, prep::ForwardDiffTwoArgJacobianPrep, - ::AutoForwardDiff, + backend::AutoForwardDiff, x, contexts::Vararg{Constant,C}, ) where {F,C} fc! = with_contexts(f!, contexts...) - return jacobian(fc!, y, x, prep.config) + CHK = tag_type(backend) === Nothing + return jacobian(fc!, y, x, prep.config, Val(CHK)) end function DI.jacobian!( @@ -351,10 +358,11 @@ function DI.jacobian!( y, jac, prep::ForwardDiffTwoArgJacobianPrep, - ::AutoForwardDiff, + backend::AutoForwardDiff, x, contexts::Vararg{Constant,C}, ) where {F,C} fc! = with_contexts(f!, contexts...) - return jacobian!(jac, fc!, y, x, prep.config) + CHK = tag_type(backend) === Nothing + return jacobian!(jac, fc!, y, x, prep.config, Val(CHK)) end diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/utils.jl b/DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/utils.jl index 5cf063410..830cf8b9c 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/utils.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/utils.jl @@ -31,6 +31,7 @@ function get_tag(f::F, ::AutoForwardDiff{chunksize,Nothing}, x) where {F,chunksi return Tag(f, eltype(x)) end +tag_type(::AutoForwardDiff{chunksize,T}) where {chunksize,T} = T tag_type(f::F, backend::AutoForwardDiff, x) where {F} = typeof(get_tag(f, backend, x)) function make_dual_similar(::Type{T}, x::Number, tx::NTuple{B}) where {T,B} diff --git a/DifferentiationInterface/test/Back/ForwardDiff/test.jl b/DifferentiationInterface/test/Back/ForwardDiff/test.jl index fa8556585..e2df26c9e 100644 --- a/DifferentiationInterface/test/Back/ForwardDiff/test.jl +++ b/DifferentiationInterface/test/Back/ForwardDiff/test.jl @@ -10,8 +10,12 @@ using Test LOGGING = get(ENV, "CI", "false") == "false" +struct MyTag end + backends = [ - AutoForwardDiff(), AutoForwardDiff(; tag=:hello), AutoForwardDiff(; chunksize=5) + AutoForwardDiff(), + AutoForwardDiff(; chunksize=5), + AutoForwardDiff(; tag=ForwardDiff.Tag(MyTag(), Float64)), ] for backend in backends @@ -54,10 +58,10 @@ test_differentiation( ## Sparse -test_differentiation(MyAutoSparse.(backends[1:2]), default_scenarios(); logging=LOGGING); +test_differentiation(MyAutoSparse(AutoForwardDiff()), default_scenarios(); logging=LOGGING); test_differentiation( - MyAutoSparse.(backends[1:2]), + MyAutoSparse(AutoForwardDiff()), sparse_scenarios(; include_constantified=true); sparsity=true, logging=LOGGING, @@ -78,7 +82,9 @@ test_differentiation(AutoForwardDiff(), static_scenarios(); logging=LOGGING) excluded=[:hessian, :pullback], # TODO: figure this out logging=LOGGING, ) - @testset "$(row[:scenario])" for row in eachrow(data) - @test row[:allocs] == 0 + @testset "Analyzing benchmark results" begin + @testset "$(row[:scenario])" for row in eachrow(data) + @test row[:allocs] == 0 + end end end;