Tag depends on ForwardDiff backend object (#229)

* Tag depends on ForwardDiff backend object

* Remove map in make_dual

Author: gdalle

DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/DifferentiationInterfaceForwardDiffExt.jl

@@ -1,6 +1,7 @@
 module DifferentiationInterfaceForwardDiffExt
 
 using ADTypes: AbstractADType, AutoForwardDiff
+using Base: Fix1
 import DifferentiationInterface as DI
 using DifferentiationInterface:
     DerivativeExtras,

DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/onearg.jl

@@ -4,8 +4,8 @@ struct ForwardDiffOneArgPushforwardExtras{T,X} <: PushforwardExtras
     xdual_tmp::X
 end
 
-function DI.prepare_pushforward(f, ::AutoForwardDiff, x, dx)
-    T = tag_type(f, x)
+function DI.prepare_pushforward(f, backend::AutoForwardDiff, x, dx)
+    T = tag_type(f, backend, x)
     xdual_tmp = make_dual(T, x, dx)
     return ForwardDiffOneArgPushforwardExtras{T,typeof(xdual_tmp)}(xdual_tmp)
 end

DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/twoarg.jl

@@ -5,8 +5,8 @@ struct ForwardDiffTwoArgPushforwardExtras{T,X,Y} <: PushforwardExtras
     ydual_tmp::Y
 end
 
-function DI.prepare_pushforward(f!, y, ::AutoForwardDiff, x, dx)
-    T = tag_type(f!, x)
+function DI.prepare_pushforward(f!, y, backend::AutoForwardDiff, x, dx)
+    T = tag_type(f!, backend, x)
     xdual_tmp = make_dual(T, x, dx)
     ydual_tmp = make_dual(T, y, similar(y))
     return ForwardDiffTwoArgPushforwardExtras{T,typeof(xdual_tmp),typeof(ydual_tmp)}(

DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/utils.jl

@@ -1,41 +1,22 @@
 choose_chunk(::AutoForwardDiff{nothing}, x) = Chunk(x)
 choose_chunk(::AutoForwardDiff{C}, x) where {C} = Chunk{C}()
 
-tag_type(::F, x::Number) where {F} = Tag{F,typeof(x)}
-tag_type(::F, x::AbstractArray) where {F} = Tag{F,eltype(x)}
+tag_type(f, ::AutoForwardDiff{C,T}, x) where {C,T} = T
+tag_type(f, ::AutoForwardDiff{C,Nothing}, x) where {C} = Tag{typeof(f),eltype(x)}
 
 make_dual(::Type{T}, x::Number, dx) where {T} = Dual{T}(x, dx)
-make_dual(::Type{T}, x::AbstractArray, dx) where {T} = Dual{T}.(x, dx)
+make_dual(::Type{T}, x, dx) where {T} = Dual{T}.(x, dx)  # TODO: map causes Enzyme to fail
 
-function make_dual!(::Type{T}, xdual, x::AbstractArray, dx) where {T}
-    for i in eachindex(xdual, x, dx)
-        xdual[i] = Dual{T}(x[i], dx[i])
-    end
-    return nothing
-end
+make_dual!(::Type{T}, xdual, x, dx) where {T} = map!(Dual{T}, xdual, x, dx)
 
 myvalue(::Type{T}, ydual::Number) where {T} = value(T, ydual)
-myvalue(::Type{T}, ydual::AbstractArray) where {T} = value.(T, ydual)
+myvalue(::Type{T}, ydual) where {T} = map(Fix1(value, T), ydual)
 
-function myvalue!(::Type{T}, y::AbstractArray, ydual) where {T}
-    for i in eachindex(y, ydual)
-        y[i] = value(T, ydual[i])
-    end
-    return nothing
-end
+myvalue!(::Type{T}, y, ydual) where {T} = map!(Fix1(value, T), y, ydual)
 
 myderivative(::Type{T}, ydual::Number) where {T} = extract_derivative(T, ydual)
-myderivative(::Type{T}, ydual::AbstractArray) where {T} = extract_derivative(T, ydual)
-
-function myderivative!(::Type{T}, dy, ydual::AbstractArray) where {T}
-    extract_derivative!(T, dy, ydual)
-    return nothing
-end
+myderivative(::Type{T}, ydual) where {T} = map(Fix1(extract_derivative, T), ydual)
 
-function myvalueandderivative!(::Type{T}, y, dy, ydual::AbstractArray) where {T}
-    for i in eachindex(y, dy, ydual)
-        y[i] = value(T, ydual[i])
-        dy[i] = extract_derivative(T, ydual[i])
-    end
-    return nothing
+function myderivative!(::Type{T}, dy, ydual) where {T}
+    return map!(Fix1(extract_derivative, T), dy, ydual)
 end