Batched ForwardDiff pushforward (#328)

* Batched ForwardDiff pushforward

* Ambiguity

* Typo

* Add batched mode to AutoFromPrimitive

* Remove type test

* Fix type instability

Author: gdalle

DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/DifferentiationInterfaceForwardDiffExt.jl

@@ -1,10 +1,11 @@
 module DifferentiationInterfaceForwardDiffExt
 
 using ADTypes: AbstractADType, AutoForwardDiff
-using Base: Fix1
+using Base: Fix1, Fix2
 using Compat
 import DifferentiationInterface as DI
 using DifferentiationInterface:
+    Batch,
     DerivativeExtras,
     GradientExtras,
     HessianExtras,
@@ -33,6 +34,8 @@ using ForwardDiff:
     hessian!,
     jacobian,
     jacobian!,
+    npartials,
+    partials,
     value
 using LinearAlgebra: dot, mul!
 

DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/onearg.jl

@@ -6,7 +6,15 @@ end
 
 function DI.prepare_pushforward(f::F, backend::AutoForwardDiff, x, dx) where {F}
     T = tag_type(f, backend, x)
-    xdual_tmp = make_dual_similar(T, x)
+    xdual_tmp = make_dual_similar(T, x, dx)
+    return ForwardDiffOneArgPushforwardExtras{T,typeof(xdual_tmp)}(xdual_tmp)
+end
+
+function DI.prepare_pushforward_batched(
+    f::F, backend::AutoForwardDiff, x, dx::Batch
+) where {F}
+    T = tag_type(f, backend, x)
+    xdual_tmp = make_dual_similar(T, x, dx)
     return ForwardDiffOneArgPushforwardExtras{T,typeof(xdual_tmp)}(xdual_tmp)
 end
 
@@ -61,56 +69,25 @@ function DI.pushforward!(
     return dy
 end
 
-## Second derivative
-
-function DI.prepare_second_derivative(f::F, backend::AutoForwardDiff, x) where {F}
-    return NoSecondDerivativeExtras()
-end
-
-function DI.second_derivative(
-    f::F, backend::AutoForwardDiff, x, ::NoSecondDerivativeExtras
-) where {F}
-    T = tag_type(f, backend, x)
-    xdual = make_dual(T, x, one(x))
-    T2 = tag_type(f, backend, xdual)
-    ydual = f(make_dual(T2, xdual, one(xdual)))
-    return myderivative(T, myderivative(T2, ydual))
-end
-
-function DI.second_derivative!(
-    f::F, der2, backend::AutoForwardDiff, x, ::NoSecondDerivativeExtras
-) where {F}
-    T = tag_type(f, backend, x)
-    xdual = make_dual(T, x, one(x))
-    T2 = tag_type(f, backend, xdual)
-    ydual = f(make_dual(T2, xdual, one(xdual)))
-    return myderivative!(T, der2, myderivative(T2, ydual))
-end
-
-function DI.value_derivative_and_second_derivative(
-    f::F, backend::AutoForwardDiff, x, ::NoSecondDerivativeExtras
-) where {F}
-    T = tag_type(f, backend, x)
-    xdual = make_dual(T, x, one(x))
-    T2 = tag_type(f, backend, xdual)
-    ydual = f(make_dual(T2, xdual, one(xdual)))
-    y = myvalue(T, myvalue(T2, ydual))
-    der = myderivative(T, myvalue(T2, ydual))
-    der2 = myderivative(T, myderivative(T2, ydual))
-    return y, der, der2
+function DI.pushforward_batched(
+    f::F, ::AutoForwardDiff, x, dx::Batch{B}, extras::ForwardDiffOneArgPushforwardExtras{T}
+) where {F,T,B}
+    ydual = compute_ydual_onearg(f, x, dx, extras)
+    new_dy = mypartials(T, Val(B), ydual)
+    return new_dy
 end
 
-function DI.value_derivative_and_second_derivative!(
-    f::F, der, der2, backend::AutoForwardDiff, x, ::NoSecondDerivativeExtras
-) where {F}
-    T = tag_type(f, backend, x)
-    xdual = make_dual(T, x, one(x))
-    T2 = tag_type(f, backend, xdual)
-    ydual = f(make_dual(T2, xdual, one(xdual)))
-    y = myvalue(T, myvalue(T2, ydual))
-    myderivative!(T, der, myvalue(T2, ydual))
-    myderivative!(T, der2, myderivative(T2, ydual))
-    return y, der, der2
+function DI.pushforward_batched!(
+    f::F,
+    dy::Batch{B},
+    ::AutoForwardDiff,
+    x,
+    dx::Batch{B},
+    extras::ForwardDiffOneArgPushforwardExtras{T},
+) where {F,T,B}
+    ydual = compute_ydual_onearg(f, x, dx, extras)
+    mypartials!(T, dy, ydual)
+    return dy
 end
 
 ## Gradient
@@ -188,6 +165,58 @@ function DI.jacobian(
     return jacobian(f, x, extras.config)
 end
 
+## Second derivative
+
+function DI.prepare_second_derivative(f::F, backend::AutoForwardDiff, x) where {F}
+    return NoSecondDerivativeExtras()
+end
+
+function DI.second_derivative(
+    f::F, backend::AutoForwardDiff, x, ::NoSecondDerivativeExtras
+) where {F}
+    T = tag_type(f, backend, x)
+    xdual = make_dual(T, x, one(x))
+    T2 = tag_type(f, backend, xdual)
+    ydual = f(make_dual(T2, xdual, one(xdual)))
+    return myderivative(T, myderivative(T2, ydual))
+end
+
+function DI.second_derivative!(
+    f::F, der2, backend::AutoForwardDiff, x, ::NoSecondDerivativeExtras
+) where {F}
+    T = tag_type(f, backend, x)
+    xdual = make_dual(T, x, one(x))
+    T2 = tag_type(f, backend, xdual)
+    ydual = f(make_dual(T2, xdual, one(xdual)))
+    return myderivative!(T, der2, myderivative(T2, ydual))
+end
+
+function DI.value_derivative_and_second_derivative(
+    f::F, backend::AutoForwardDiff, x, ::NoSecondDerivativeExtras
+) where {F}
+    T = tag_type(f, backend, x)
+    xdual = make_dual(T, x, one(x))
+    T2 = tag_type(f, backend, xdual)
+    ydual = f(make_dual(T2, xdual, one(xdual)))
+    y = myvalue(T, myvalue(T2, ydual))
+    der = myderivative(T, myvalue(T2, ydual))
+    der2 = myderivative(T, myderivative(T2, ydual))
+    return y, der, der2
+end
+
+function DI.value_derivative_and_second_derivative!(
+    f::F, der, der2, backend::AutoForwardDiff, x, ::NoSecondDerivativeExtras
+) where {F}
+    T = tag_type(f, backend, x)
+    xdual = make_dual(T, x, one(x))
+    T2 = tag_type(f, backend, xdual)
+    ydual = f(make_dual(T2, xdual, one(xdual)))
+    y = myvalue(T, myvalue(T2, ydual))
+    myderivative!(T, der, myvalue(T2, ydual))
+    myderivative!(T, der2, myderivative(T2, ydual))
+    return y, der, der2
+end
+
 ## Hessian
 
 struct ForwardDiffHessianExtras{C1,C2,C3} <: HessianExtras

DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/twoarg.jl

@@ -7,8 +7,19 @@ end
 
 function DI.prepare_pushforward(f!::F, y, backend::AutoForwardDiff, x, dx) where {F}
     T = tag_type(f!, backend, x)
-    xdual_tmp = make_dual_similar(T, x)
-    ydual_tmp = make_dual_similar(T, y)
+    xdual_tmp = make_dual_similar(T, x, dx)
+    ydual_tmp = make_dual_similar(T, y, dx)  # dx only for batch size
+    return ForwardDiffTwoArgPushforwardExtras{T,typeof(xdual_tmp),typeof(ydual_tmp)}(
+        xdual_tmp, ydual_tmp
+    )
+end
+
+function DI.prepare_pushforward_batched(
+    f!::F, y, backend::AutoForwardDiff, x, dx::Batch
+) where {F}
+    T = tag_type(f!, backend, x)
+    xdual_tmp = make_dual_similar(T, x, dx)
+    ydual_tmp = make_dual_similar(T, y, dx)  # dx only for batch size
     return ForwardDiffTwoArgPushforwardExtras{T,typeof(xdual_tmp),typeof(ydual_tmp)}(
         xdual_tmp, ydual_tmp
     )
@@ -66,6 +77,33 @@ function DI.pushforward!(
     return dy
 end
 
+function DI.pushforward_batched(
+    f!::F,
+    y,
+    ::AutoForwardDiff,
+    x,
+    dx::Batch{B},
+    extras::ForwardDiffTwoArgPushforwardExtras{T},
+) where {F,T,B}
+    ydual_tmp = compute_ydual_twoarg(f!, y, x, dx, extras)
+    dy = mypartials(T, Val(B), ydual_tmp)
+    return dy
+end
+
+function DI.pushforward_batched!(
+    f!::F,
+    y,
+    dy::Batch{B},
+    ::AutoForwardDiff,
+    x,
+    dx::Batch{B},
+    extras::ForwardDiffTwoArgPushforwardExtras{T},
+) where {F,T,B}
+    ydual_tmp = compute_ydual_twoarg(f!, y, x, dx, extras)
+    mypartials!(T, dy, ydual_tmp)
+    return dy
+end
+
 ## Derivative
 
 struct ForwardDiffTwoArgDerivativeExtras{C} <: DerivativeExtras

DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/utils.jl

@@ -1,23 +1,47 @@
 choose_chunk(::AutoForwardDiff{nothing}, x) = Chunk(x)
-choose_chunk(::AutoForwardDiff{C}, x) where {C} = Chunk{C}()
+choose_chunk(::AutoForwardDiff{C}, x) where {C} = Chunk{min(length(x), C)}()
 
 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_similar(::Type{T}, x::Number) where {T} = Dual{T}(x, x)
-make_dual_similar(::Type{T}, x) where {T} = similar(x, Dual{T,eltype(x),1})
+make_dual_similar(::Type{T}, x::Number, dx::Number) where {T} = Dual{T}(x, dx)
+make_dual_similar(::Type{T}, x, dx) where {T} = similar(x, Dual{T,eltype(x),1})
 
-make_dual(::Type{T}, x::Number, dx) where {T} = Dual{T}(x, dx)
+function make_dual_similar(::Type{T}, x, dx::Batch{B}) where {T,B}
+    return similar(x, Dual{T,eltype(x),B})
+end
+
+make_dual(::Type{T}, x::Number, dx::Number) where {T} = Dual{T}(x, dx)
 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) where {T} = map(Fix1(value, T), ydual)
+function make_dual!(::Type{T}, xdual, x, dx::Batch{B}) where {T,B}
+    return map!(Dual{T}, xdual, x, dx.elements...)
+end
 
-myvalue!(::Type{T}, y, ydual) where {T} = map!(Fix1(value, T), y, ydual)
+myvalue(::Type{T}, ydual::Dual{T}) where {T} = value(T, ydual)
+myvalue(::Type{T}, ydual) where {T} = map(Fix1(myvalue, T), ydual)
+
+function myvalue!(::Type{T}, y, ydual) where {T}
+    return map!(Fix1(myvalue, T), y, ydual)
+end
 
-myderivative(::Type{T}, ydual::Number) where {T} = extract_derivative(T, ydual)
-myderivative(::Type{T}, ydual) where {T} = map(Fix1(extract_derivative, T), ydual)
+myderivative(::Type{T}, ydual::Dual{T}) where {T} = extract_derivative(T, ydual)
+myderivative(::Type{T}, ydual) where {T} = map(Fix1(myderivative, T), ydual)
 
 function myderivative!(::Type{T}, dy, ydual) where {T}
-    return map!(Fix1(extract_derivative, T), dy, ydual)
+    return map!(Fix1(myderivative, T), dy, ydual)
+end
+
+function mypartials(::Type{T}, ::Val{B}, ydual) where {T,B}
+    elements = ntuple(Val(B)) do b
+        partials.(T, ydual, b)
+    end
+    return Batch(elements)
+end
+
+function mypartials!(::Type{T}, dy::Batch{B}, ydual) where {T,B}
+    for b in eachindex(dy.elements)
+        dy.elements[b] .= partials.(T, ydual, b)
+    end
+    return dy
 end