Test batched operators (#331)

* Add batched tests

* Full batched interface with tests

* Fixes

* Fix

* Don't batchify by default

* Coverage in DIT

* Typos and coverage

Author: gdalle

DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/utils.jl

@@ -7,29 +7,41 @@ tag_type(f, ::AutoForwardDiff{C,Nothing}, x) where {C} = Tag{typeof(f),eltype(x)
 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})
 
+function make_dual_similar(::Type{T}, x::Number, dx::Batch{B,<:Number}) where {T,B}
+    return Dual{T}(x, dx.elements)
+end
+
 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)
+function make_dual(::Type{T}, x::Number, dx::Number) where {T}
+    return Dual{T}(x, dx)
+end
 
-function make_dual!(::Type{T}, xdual, x, dx::Batch{B}) where {T,B}
-    return map!(Dual{T}, xdual, x, dx.elements...)
+function make_dual(::Type{T}, x::Number, dx::Batch{B,<:Number}) where {T,B}
+    return Dual{T}(x, dx.elements...)
 end
 
-myvalue(::Type{T}, ydual::Dual{T}) where {T} = value(T, ydual)
-myvalue(::Type{T}, ydual) where {T} = map(Fix1(myvalue, T), ydual)
+function make_dual!(::Type{T}, xdual, x, dx) where {T}
+    return xdual .= Dual{T}.(x, dx)
+end
 
-function myvalue!(::Type{T}, y, ydual) where {T}
-    return map!(Fix1(myvalue, T), y, ydual)
+function make_dual!(::Type{T}, xdual, x, dx::Batch{B}) where {T,B}
+    return xdual .= Dual{T}.(x, dx.elements...)
 end
 
+myvalue(::Type{T}, ydual::Dual{T}) where {T} = value(T, ydual)
+myvalue(::Type{T}, ydual) where {T} = myvalue.(T, ydual)
+myvalue!(::Type{T}, y, ydual) where {T} = y .= myvalue.(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)
+myderivative(::Type{T}, ydual) where {T} = myderivative.(T, ydual)
+myderivative!(::Type{T}, dy, ydual) where {T} = dy .= myderivative.(T, ydual)
 
-function myderivative!(::Type{T}, dy, ydual) where {T}
-    return map!(Fix1(myderivative, T), dy, ydual)
+function mypartials(::Type{T}, ::Val{B}, ydual::Dual) where {T,B}
+    elements = partials(T, ydual).values
+    return Batch(elements)
 end
 
 function mypartials(::Type{T}, ::Val{B}, ydual) where {T,B}

DifferentiationInterface/src/DifferentiationInterface.jl

@@ -56,13 +56,16 @@ include("utils/exceptions.jl")
 include("utils/maybe.jl")
 
 include("first_order/pushforward.jl")
+include("first_order/pushforward_batched.jl")
 include("first_order/pullback.jl")
+include("first_order/pullback_batched.jl")
 include("first_order/derivative.jl")
 include("first_order/gradient.jl")
 include("first_order/jacobian.jl")
 
 include("second_order/second_derivative.jl")
 include("second_order/hvp.jl")
+include("second_order/hvp_batched.jl")
 include("second_order/hessian.jl")
 
 include("sparse/fallbacks.jl")

DifferentiationInterface/src/first_order/derivative.jl

@@ -72,6 +72,8 @@ end
 
 ## One argument
 
+### Without extras
+
 function value_and_derivative(f::F, backend::AbstractADType, x) where {F}
     return value_and_derivative(f, backend, x, prepare_derivative(f, backend, x))
 end
@@ -88,6 +90,8 @@ function derivative!(f::F, der, backend::AbstractADType, x) where {F}
     return derivative!(f, der, backend, x, prepare_derivative(f, backend, x))
 end
 
+### With extras
+
 function value_and_derivative(
     f::F, backend::AbstractADType, x, extras::PushforwardDerivativeExtras
 ) where {F}
@@ -114,6 +118,8 @@ end
 
 ## Two arguments
 
+### Without extras
+
 function value_and_derivative(f!::F, y, backend::AbstractADType, x) where {F}
     return value_and_derivative(f!, y, backend, x, prepare_derivative(f!, y, backend, x))
 end
@@ -132,6 +138,8 @@ function derivative!(f!::F, y, der, backend::AbstractADType, x) where {F}
     return derivative!(f!, y, der, backend, x, prepare_derivative(f!, y, backend, x))
 end
 
+### With extras
+
 function value_and_derivative(
     f!::F, y, backend::AbstractADType, x, extras::PushforwardDerivativeExtras
 ) where {F}

DifferentiationInterface/src/first_order/gradient.jl

@@ -62,6 +62,8 @@ end
 
 ## One argument
 
+### Without extras
+
 function value_and_gradient(f::F, backend::AbstractADType, x) where {F}
     return value_and_gradient(f, backend, x, prepare_gradient(f, backend, x))
 end
@@ -78,6 +80,8 @@ function gradient!(f::F, der, backend::AbstractADType, x) where {F}
     return gradient!(f, der, backend, x, prepare_gradient(f, backend, x))
 end
 
+### With extras
+
 function value_and_gradient(
     f::F, backend::AbstractADType, x, extras::PullbackGradientExtras
 ) where {F}

DifferentiationInterface/src/first_order/jacobian.jl

@@ -104,6 +104,8 @@ end
 
 ## One argument
 
+### Without extras
+
 function jacobian(f::F, backend::AbstractADType, x) where {F}
     return jacobian(f, backend, x, prepare_jacobian(f, backend, x))
 end
@@ -120,6 +122,8 @@ function value_and_jacobian!(f::F, jac, backend::AbstractADType, x) where {F}
     return value_and_jacobian!(f, jac, backend, x, prepare_jacobian(f, backend, x))
 end
 
+### With extras
+
 function jacobian(f::F, backend::AbstractADType, x, extras::JacobianExtras) where {F}
     return jacobian_aux((f,), backend, x, extras)
 end
@@ -142,6 +146,8 @@ end
 
 ## Two arguments
 
+### Without extras
+
 function jacobian(f!::F, y, backend::AbstractADType, x) where {F}
     return jacobian(f!, y, backend, x, prepare_jacobian(f!, y, backend, x))
 end
@@ -158,6 +164,8 @@ function value_and_jacobian!(f!::F, y, jac, backend::AbstractADType, x) where {F
     return value_and_jacobian!(f!, y, jac, backend, x, prepare_jacobian(f!, y, backend, x))
 end
 
+### With extras
+
 function jacobian(f!::F, y, backend::AbstractADType, x, extras::JacobianExtras) where {F}
     return jacobian_aux((f!, y), backend, x, extras)
 end

DifferentiationInterface/src/first_order/pullback.jl

@@ -12,8 +12,6 @@ Create an `extras` object that can be given to [`pullback`](@ref) and its varian
 """
 function prepare_pullback end
 
-function prepare_pullback_batched end
-
 """
     prepare_pullback_same_point(f,     backend, x, dy) -> extras_same
     prepare_pullback_same_point(f!, y, backend, x, dy) -> extras_same
@@ -26,8 +24,6 @@ Create an `extras_same` object that can be given to [`pullback`](@ref) and its v
 """
 function prepare_pullback_same_point end
 
-function prepare_pullback_batched_same_point end
-
 """
     value_and_pullback(f,     backend, x, dy, [extras]) -> (y, dx)
     value_and_pullback(f!, y, backend, x, dy, [extras]) -> (y, dx)
@@ -55,8 +51,6 @@ Compute the pullback of the function `f` at point `x` with seed `dy`.
 """
 function pullback end
 
-function pullback_batched end
-
 """
     pullback!(f,     dx, backend, x, dy, [extras]) -> dx
     pullback!(f!, y, dx, backend, x, dy, [extras]) -> dx
@@ -65,8 +59,6 @@ Compute the pullback of the function `f` at point `x` with seed `dy`, overwritin
 """
 function pullback! end
 
-function pullback_batched! end
-
 ## Preparation
 
 ### Extras types
@@ -84,7 +76,7 @@ struct PushforwardPullbackExtras{E} <: PullbackExtras
     pushforward_extras::E
 end
 
-## Standard
+## Different point
 
 function prepare_pullback(f::F, backend::AbstractADType, x, dy) where {F}
     return prepare_pullback_aux(f, backend, x, dy, pullback_performance(backend))
@@ -114,7 +106,7 @@ function prepare_pullback_aux(f!, y, backend, x, dy, ::PullbackFast)
     throw(MissingBackendError(backend))
 end
 
-### Standard, same point
+### Same point
 
 function prepare_pullback_same_point(
     f::F, backend::AbstractADType, x, dy, extras::PullbackExtras
@@ -138,33 +130,9 @@ function prepare_pullback_same_point(f!::F, y, backend::AbstractADType, x, dy) w
     return prepare_pullback_same_point(f!, y, backend, x, dy, extras)
 end
 
-### Batched
-
-function prepare_pullback_batched(f::F, backend::AbstractADType, x, dy::Batch) where {F}
-    return prepare_pullback(f, backend, x, first(dy.elements))
-end
-
-function prepare_pullback_batched(f!::F, y, backend::AbstractADType, x, dy::Batch) where {F}
-    return prepare_pullback(f!, y, backend, x, first(dy.elements))
-end
-
-### Batched, same point
-
-function prepare_pullback_batched_same_point(
-    f::F, backend::AbstractADType, x, dy::Batch, extras::PullbackExtras
-) where {F}
-    return prepare_pullback_same_point(f, backend, x, first(dy.elements), extras)
-end
-
-function prepare_pullback_batched_same_point(
-    f!::F, y, backend::AbstractADType, x, dy::Batch, extras::PullbackExtras
-) where {F}
-    return prepare_pullback_same_point(f!, y, backend, x, first(dy.elements), extras)
-end
-
 ## One argument
 
-### Standard
+### Without extras
 
 function value_and_pullback(f::F, backend::AbstractADType, x, dy) where {F}
     return value_and_pullback(f, backend, x, dy, prepare_pullback(f, backend, x, dy))
@@ -182,6 +150,8 @@ function pullback!(f::F, dx, backend::AbstractADType, x, dy) where {F}
     return pullback!(f, dx, backend, x, dy, prepare_pullback(f, backend, x, dy))
 end
 
+### With extras
+
 function value_and_pullback(
     f::F, backend, x, dy, extras::PushforwardPullbackExtras
 ) where {F}
@@ -220,29 +190,9 @@ function pullback!(
     return value_and_pullback!(f, dx, backend, x, dy, extras)[2]
 end
 
-### Batched
-
-function pullback_batched(
-    f::F, backend::AbstractADType, x, dy::Batch{B}, extras::PullbackExtras
-) where {F,B}
-    dx_elements = ntuple(Val(B)) do b
-        pullback(f, backend, x, dy.elements[b], extras)
-    end
-    return Batch(dx_elements)
-end
-
-function pullback_batched!(
-    f::F, dx::Batch, backend::AbstractADType, x, dy::Batch, extras::PullbackExtras
-) where {F}
-    for b in eachindex(dx.elements, dy.elements)
-        pullback!(f, dx.elements[b], backend, x, dy.elements[b], extras)
-    end
-    return dx
-end
-
 ## Two arguments
 
-### Standard
+### Without extras
 
 function value_and_pullback(f!::F, y, backend::AbstractADType, x, dy) where {F}
     return value_and_pullback(
@@ -264,6 +214,8 @@ function pullback!(f!::F, y, dx, backend::AbstractADType, x, dy) where {F}
     return pullback!(f!, y, dx, backend, x, dy, prepare_pullback(f!, y, backend, x, dy))
 end
 
+### With extras
+
 function value_and_pullback(
     f!::F, y, backend, x, dy, extras::PushforwardPullbackExtras
 ) where {F}
@@ -297,23 +249,3 @@ function pullback!(
 ) where {F}
     return value_and_pullback!(f!, y, dx, backend, x, dy, extras)[2]
 end
-
-### Batched
-
-function pullback_batched(
-    f!::F, y, backend::AbstractADType, x, dy::Batch{B}, extras::PullbackExtras
-) where {F,B}
-    dx_elements = ntuple(Val(B)) do b
-        pullback(f!, y, backend, x, dy.elements[b], extras)
-    end
-    return Batch(dx_elements)
-end
-
-function pullback_batched!(
-    f!::F, y, dx::Batch, backend::AbstractADType, x, dy::Batch, extras::PullbackExtras
-) where {F}
-    for b in eachindex(dx.elements, dy.elements)
-        pullback!(f!, y, dx.elements[b], backend, x, dy.elements[b], extras)
-    end
-    return dx
-end