Fix PolyesterForwardDiff

Author: gdalle

DifferentiationInterface/Project.toml

@@ -39,7 +39,7 @@ DifferentiationInterfaceFiniteDifferencesExt = "FiniteDifferences"
 DifferentiationInterfaceForwardDiffExt = ["ForwardDiff", "DiffResults"]
 DifferentiationInterfaceGTPSAExt = "GTPSA"
 DifferentiationInterfaceMooncakeExt = "Mooncake"
-DifferentiationInterfacePolyesterForwardDiffExt = "PolyesterForwardDiff"
+DifferentiationInterfacePolyesterForwardDiffExt = ["PolyesterForwardDiff", "ForwardDiff", "DiffResults"]
 DifferentiationInterfaceReverseDiffExt = ["ReverseDiff", "DiffResults"]
 DifferentiationInterfaceSparseArraysExt = "SparseArrays"
 DifferentiationInterfaceSparseConnectivityTracerExt = "SparseConnectivityTracer"
@@ -115,7 +115,6 @@ test = [
     "ComponentArrays",
     "DataFrames",
     "ExplicitImports",
-    "ForwardDiff",
     "JET",
     "JLArrays",
     "JuliaFormatter",

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/utils.jl

@@ -48,13 +48,13 @@ force_annotation(f::F) where {F<:Annotation} = f
 force_annotation(f::F) where {F} = Const(f)
 
 @inline function _translate(
-    ::AutoEnzyme, ::Mode, ::Val{B}, c::Union{DI.Constant,DI.BackendContext}
+    ::AutoEnzyme, ::Mode, ::Val{B}, c::DI.GeneralizedConstant
 ) where {B}
     return Const(DI.unwrap(c))
 end
 
 @inline function _translate(
-    backend::AutoEnzyme, mode::Mode, ::Val{B}, c::DI.Cache
+    ::AutoEnzyme, ::Mode, ::Val{B}, c::DI.GeneralizedCache
 ) where {B}
     if B == 1
         return Duplicated(DI.unwrap(c), make_zero(DI.unwrap(c)))

DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/misc.jl

@@ -3,27 +3,16 @@ DI.overloaded_input_type(prep::ForwardDiffOneArgPushforwardPrep) = typeof(prep.x
 DI.overloaded_input_type(prep::ForwardDiffTwoArgPushforwardPrep) = typeof(prep.xdual_tmp)
 
 function DI.overloaded_input(
-    ::typeof(DI.pushforward),
-    f::F,
-    backend::AutoForwardDiff,
-    x,
-    tx::NTuple{B},
-    contexts::Vararg{DI.Context,C},
-) where {F,B,C}
+    ::typeof(DI.pushforward), f::F, backend::AutoForwardDiff, x, tx::NTuple{B}
+) where {F,B}
     T = tag_type(f, backend, x)
     xdual = make_dual(T, x, tx)
     return xdual
 end
 
 function DI.overloaded_input(
-    ::typeof(DI.pushforward),
-    f!::F,
-    y,
-    backend::AutoForwardDiff,
-    x,
-    tx::NTuple{B},
-    contexts::Vararg{DI.Context,C},
-) where {F,B,C}
+    ::typeof(DI.pushforward), f!::F, y, backend::AutoForwardDiff, x, tx::NTuple{B}
+) where {F,B}
     T = tag_type(f!, backend, x)
     xdual = make_dual(T, x, tx)
     return xdual

DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/utils.jl

@@ -82,7 +82,11 @@ function mypartials!(::Type{T}, ty::NTuple{B}, ydual) where {T,B}
     return ty
 end
 
-_translate(::Type{D}, c::DI.GeneralizedConstant) where {D<:Dual} = DI.unwrap(c)
+function _translate(
+    ::Type{D}, c::Union{DI.GeneralizedConstant,DI.PrepContext}
+) where {D<:Dual}
+    return DI.unwrap(c)
+end
 function _translate(::Type{D}, c::DI.Cache) where {D<:Dual}
     c0 = DI.unwrap(c)
     return similar(c0, D)
@@ -95,7 +99,11 @@ function translate(::Type{D}, contexts::NTuple{C,DI.Context}) where {D<:Dual,C}
     return new_contexts
 end
 
-_translate_toprep(::Type{D}, c::DI.GeneralizedConstant) where {D<:Dual} = nothing
+function _translate_toprep(
+    ::Type{D}, c::Union{DI.GeneralizedConstant,DI.PrepContext}
+) where {D<:Dual}
+    return nothing
+end
 function _translate_toprep(::Type{D}, c::DI.Cache) where {D<:Dual}
     c0 = DI.unwrap(c)
     return similar(c0, D)
@@ -108,7 +116,7 @@ function translate_toprep(::Type{D}, contexts::NTuple{C,DI.Context}) where {D<:D
     return new_contexts
 end
 
-_translate_prepared(c::DI.GeneralizedConstant, _pc) = DI.unwrap(c)
+_translate_prepared(c::Union{DI.GeneralizedConstant,DI.PrepContext}, _pc) = DI.unwrap(c)
 _translate_prepared(_c::DI.Cache, pc) = pc
 
 function translate_prepared(

DifferentiationInterface/ext/DifferentiationInterfacePolyesterForwardDiffExt/DifferentiationInterfacePolyesterForwardDiffExt.jl

@@ -4,31 +4,22 @@ using ADTypes: AutoForwardDiff, AutoPolyesterForwardDiff
 import DifferentiationInterface as DI
 using LinearAlgebra: mul!
 using PolyesterForwardDiff: threaded_gradient!, threaded_jacobian!
-using PolyesterForwardDiff.ForwardDiff: Chunk
-using PolyesterForwardDiff.ForwardDiff.DiffResults: DiffResults
+using ForwardDiff: Chunk
+using DiffResults: DiffResults
+
+const FDExt = Base.get_extension(DI, :DifferentiationInterfaceForwardDiffExt)
+@assert !isnothing(FDExt)
 
 function single_threaded(backend::AutoPolyesterForwardDiff{chunksize,T}) where {chunksize,T}
     return AutoForwardDiff(; chunksize, tag=backend.tag)
 end
 
 DI.check_available(::AutoPolyesterForwardDiff) = true
+DI.inner_preparation_behavior(::AutoPolyesterForwardDiff) = DI.PrepareInnerOverload()
 
-function DI.pick_batchsize(backend::AutoPolyesterForwardDiff, x::AbstractArray)
-    return DI.pick_batchsize(single_threaded(backend), x)
-end
-
-function DI.pick_batchsize(backend::AutoPolyesterForwardDiff, N::Integer)
-    return DI.pick_batchsize(single_threaded(backend), N)
-end
-
-function DI.threshold_batchsize(
-    backend::AutoPolyesterForwardDiff{chunksize1}, chunksize2::Integer
-) where {chunksize1}
-    chunksize = isnothing(chunksize1) ? nothing : min(chunksize1, chunksize2)
-    return AutoPolyesterForwardDiff(; chunksize, tag=backend.tag)
-end
-
+include("utils.jl")
 include("onearg.jl")
 include("twoarg.jl")
+include("misc.jl")
 
 end # module

DifferentiationInterface/ext/DifferentiationInterfacePolyesterForwardDiffExt/misc.jl

@@ -0,0 +1,11 @@
+function DI.overloaded_input(
+    ::typeof(DI.pushforward), f::F, backend::AutoPolyesterForwardDiff, x, tx::NTuple{B}
+) where {F,B}
+    return DI.overloaded_input(DI.pushforward, f, single_threaded(backend), x, tx)
+end
+
+function DI.overloaded_input(
+    ::typeof(DI.pushforward), f!::F, y, backend::AutoPolyesterForwardDiff, x, tx::NTuple{B}
+) where {F,B}
+    return DI.overloaded_input(DI.pushforward, f!, y, single_threaded(backend), x, tx)
+end

DifferentiationInterface/ext/DifferentiationInterfacePolyesterForwardDiffExt/onearg.jl

@@ -107,45 +107,61 @@ end
 
 ## Gradient
 
-struct PolyesterForwardDiffGradientPrep{chunksize} <: DI.GradientPrep
+struct PolyesterForwardDiffGradientPrep{chunksize,P} <: DI.GradientPrep
     chunk::Chunk{chunksize}
+    single_threaded_prep::P
 end
 
 function DI.prepare_gradient(
-    f, ::AutoPolyesterForwardDiff{chunksize}, x, contexts::Vararg{DI.Context,C}
+    f, backend::AutoPolyesterForwardDiff{chunksize}, x, contexts::Vararg{DI.Context,C}
 ) where {chunksize,C}
     if isnothing(chunksize)
         chunk = Chunk(x)
     else
         chunk = Chunk{chunksize}()
     end
-    return PolyesterForwardDiffGradientPrep(chunk)
+    single_threaded_prep = DI.prepare_gradient(f, single_threaded(backend), x, contexts...)
+    return PolyesterForwardDiffGradientPrep(chunk, single_threaded_prep)
 end
 
 function DI.value_and_gradient!(
     f,
     grad,
     prep::PolyesterForwardDiffGradientPrep,
-    ::AutoPolyesterForwardDiff,
+    backend::AutoPolyesterForwardDiff,
     x,
     contexts::Vararg{DI.Context,C},
 ) where {C}
-    fc = DI.with_contexts(f, contexts...)
-    threaded_gradient!(fc, grad, x, prep.chunk)
-    return fc(x), grad
+    if contexts isa NTuple{C,DI.GeneralizedConstant}
+        fc = DI.with_contexts(f, contexts...)
+        threaded_gradient!(fc, grad, x, prep.chunk)
+        return fc(x), grad
+    else
+        # TODO: optimize
+        return DI.value_and_gradient!(
+            f, grad, prep.single_threaded_prep, single_threaded(backend), x, contexts...
+        )
+    end
 end
 
 function DI.gradient!(
     f,
     grad,
     prep::PolyesterForwardDiffGradientPrep,
-    ::AutoPolyesterForwardDiff,
+    backend::AutoPolyesterForwardDiff,
     x,
     contexts::Vararg{DI.Context,C},
 ) where {C}
-    fc = DI.with_contexts(f, contexts...)
-    threaded_gradient!(fc, grad, x, prep.chunk)
-    return grad
+    if contexts isa NTuple{C,DI.GeneralizedConstant}
+        fc = DI.with_contexts(f, contexts...)
+        threaded_gradient!(fc, grad, x, prep.chunk)
+        return grad
+    else
+        # TODO: optimize
+        return DI.gradient!(
+            f, grad, prep.single_threaded_prep, single_threaded(backend), x, contexts...
+        )
+    end
 end
 
 function DI.value_and_gradient(
@@ -170,43 +186,57 @@ end
 
 ## Jacobian
 
-struct PolyesterForwardDiffOneArgJacobianPrep{chunksize} <: DI.JacobianPrep
+struct PolyesterForwardDiffOneArgJacobianPrep{chunksize,P} <: DI.JacobianPrep
     chunk::Chunk{chunksize}
+    single_threaded_prep::P
 end
 
 function DI.prepare_jacobian(
-    f, ::AutoPolyesterForwardDiff{chunksize}, x, contexts::Vararg{DI.Context,C}
+    f, backend::AutoPolyesterForwardDiff{chunksize}, x, contexts::Vararg{DI.Context,C}
 ) where {chunksize,C}
     if isnothing(chunksize)
         chunk = Chunk(x)
     else
         chunk = Chunk{chunksize}()
     end
-    return PolyesterForwardDiffOneArgJacobianPrep(chunk)
+    single_threaded_prep = DI.prepare_jacobian(f, single_threaded(backend), x, contexts...)
+    return PolyesterForwardDiffOneArgJacobianPrep(chunk, single_threaded_prep)
 end
 
 function DI.value_and_jacobian!(
     f,
     jac,
     prep::PolyesterForwardDiffOneArgJacobianPrep,
-    ::AutoPolyesterForwardDiff,
+    backend::AutoPolyesterForwardDiff,
     x,
     contexts::Vararg{DI.Context,C},
 ) where {C}
-    fc = DI.with_contexts(f, contexts...)
-    return fc(x), threaded_jacobian!(fc, jac, x, prep.chunk)
+    if contexts isa NTuple{C,DI.GeneralizedConstant}
+        fc = DI.with_contexts(f, contexts...)
+        return fc(x), threaded_jacobian!(fc, jac, x, prep.chunk)
+    else
+        return DI.value_and_jacobian!(
+            f, jac, prep.single_threaded_prep, single_threaded(backend), x, contexts...
+        )
+    end
 end
 
 function DI.jacobian!(
     f,
     jac,
     prep::PolyesterForwardDiffOneArgJacobianPrep,
-    ::AutoPolyesterForwardDiff,
+    backend::AutoPolyesterForwardDiff,
     x,
     contexts::Vararg{DI.Context,C},
 ) where {C}
-    fc = DI.with_contexts(f, contexts...)
-    return threaded_jacobian!(fc, jac, x, prep.chunk)
+    if contexts isa NTuple{C,DI.GeneralizedConstant}
+        fc = DI.with_contexts(f, contexts...)
+        return threaded_jacobian!(fc, jac, x, prep.chunk)
+    else
+        return DI.jacobian!(
+            f, jac, prep.single_threaded_prep, single_threaded(backend), x, contexts...
+        )
+    end
 end
 
 function DI.value_and_jacobian(
@@ -217,9 +247,8 @@ function DI.value_and_jacobian(
     contexts::Vararg{DI.Context,C},
 ) where {C}
     y = f(x, map(DI.unwrap, contexts)...)
-    return DI.value_and_jacobian!(
-        f, similar(y, length(y), length(x)), prep, backend, x, contexts...
-    )
+    jac = similar(y, length(y), length(x))
+    return DI.value_and_jacobian!(f, jac, prep, backend, x, contexts...)
 end
 
 function DI.jacobian(
@@ -230,7 +259,8 @@ function DI.jacobian(
     contexts::Vararg{DI.Context,C},
 ) where {C}
     y = f(x, map(DI.unwrap, contexts)...)
-    return DI.jacobian!(f, similar(y, length(y), length(x)), prep, backend, x, contexts...)
+    jac = similar(y, length(y), length(x))
+    return DI.jacobian!(f, jac, prep, backend, x, contexts...)
 end
 
 ## Hessian
@@ -299,7 +329,7 @@ end
 
 function DI.hvp(
     f,
-    prep::DI.HVPPrep,
+    prep::DI.ForwardOverAnythingHVPPrep,
     backend::AutoPolyesterForwardDiff,
     x,
     tx::NTuple,
@@ -313,7 +343,7 @@ end
 function DI.hvp!(
     f,
     tg::NTuple,
-    prep::DI.HVPPrep,
+    prep::DI.ForwardOverAnythingHVPPrep,
     backend::AutoPolyesterForwardDiff,
     x,
     tx::NTuple,
@@ -326,7 +356,7 @@ end
 
 function DI.gradient_and_hvp(
     f,
-    prep::DI.HVPPrep,
+    prep::DI.ForwardOverAnythingHVPPrep,
     backend::AutoPolyesterForwardDiff,
     x,
     tx::NTuple,
@@ -341,7 +371,7 @@ function DI.gradient_and_hvp!(
     f,
     grad,
     tg::NTuple,
-    prep::DI.HVPPrep,
+    prep::DI.ForwardOverAnythingHVPPrep,
     backend::AutoPolyesterForwardDiff,
     x,
     tx::NTuple,