Batched pushforward, pullback and hvp (#320)

* Batched pushforward, pullback and hvp

* Fixes, add FromPrimitive for testing

* Typo

* Typo

* Typos

* More formatting

* Reduce code duplication

* Typos

* Better display

* Typos

* Uncomment

* Printing

* Typos

* Type stability

* Typo

* Typo

* Typo

* Log Zygote

* Forward-over-reverse HVP batched for Zygote

* Typo and coverage

* Chunksize

* Funny chunk size

Author: gdalle

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/DifferentiationInterfaceEnzymeExt.jl

@@ -13,7 +13,7 @@ using DifferentiationInterface:
     NoJacobianExtras,
     NoPullbackExtras,
     NoPushforwardExtras,
-    pick_chunksize
+    pick_batchsize
 using DocStringExtensions
 using Enzyme:
     Active,

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/forward_onearg.jl

@@ -57,20 +57,20 @@ end
 
 ## Gradient
 
-struct EnzymeForwardGradientExtras{C,O} <: GradientExtras
+struct EnzymeForwardGradientExtras{B,O} <: GradientExtras
     shadow::O
 end
 
-function DI.prepare_gradient(f, ::AutoEnzyme{<:ForwardMode}, x)
-    C = pick_chunksize(length(x))
-    shadow = chunkedonehot(x, Val(C))
-    return EnzymeForwardGradientExtras{C,typeof(shadow)}(shadow)
+function DI.prepare_gradient(f, backend::AutoEnzyme{<:ForwardMode}, x)
+    B = pick_batchsize(backend, length(x))
+    shadow = chunkedonehot(x, Val(B))
+    return EnzymeForwardGradientExtras{B,typeof(shadow)}(shadow)
 end
 
 function DI.gradient(
-    f, backend::AutoEnzyme{<:ForwardMode}, x, extras::EnzymeForwardGradientExtras{C}
-) where {C}
-    grad_tup = gradient(forward_mode(backend), f, x, Val{C}(); shadow=extras.shadow)
+    f, backend::AutoEnzyme{<:ForwardMode}, x, extras::EnzymeForwardGradientExtras{B}
+) where {B}
+    grad_tup = gradient(forward_mode(backend), f, x, Val(B); shadow=extras.shadow)
     return reshape(collect(grad_tup), size(x))
 end
 
@@ -81,38 +81,38 @@ function DI.value_and_gradient(
 end
 
 function DI.gradient!(
-    f, grad, backend::AutoEnzyme{<:ForwardMode}, x, extras::EnzymeForwardGradientExtras{C}
-) where {C}
-    grad_tup = gradient(forward_mode(backend), f, x, Val{C}(); shadow=extras.shadow)
+    f, grad, backend::AutoEnzyme{<:ForwardMode}, x, extras::EnzymeForwardGradientExtras{B}
+) where {B}
+    grad_tup = gradient(forward_mode(backend), f, x, Val(B); shadow=extras.shadow)
     return copyto!(grad, grad_tup)
 end
 
 function DI.value_and_gradient!(
-    f, grad, backend::AutoEnzyme{<:ForwardMode}, x, extras::EnzymeForwardGradientExtras{C}
-) where {C}
-    grad_tup = gradient(forward_mode(backend), f, x, Val{C}(); shadow=extras.shadow)
+    f, grad, backend::AutoEnzyme{<:ForwardMode}, x, extras::EnzymeForwardGradientExtras{B}
+) where {B}
+    grad_tup = gradient(forward_mode(backend), f, x, Val(B); shadow=extras.shadow)
     return f(x), copyto!(grad, grad_tup)
 end
 
 ## Jacobian
 
-struct EnzymeForwardOneArgJacobianExtras{C,O} <: JacobianExtras
+struct EnzymeForwardOneArgJacobianExtras{B,O} <: JacobianExtras
     shadow::O
 end
 
-function DI.prepare_jacobian(f, ::AutoEnzyme{<:Union{ForwardMode,Nothing}}, x)
-    C = pick_chunksize(length(x))
-    shadow = chunkedonehot(x, Val(C))
-    return EnzymeForwardOneArgJacobianExtras{C,typeof(shadow)}(shadow)
+function DI.prepare_jacobian(f, backend::AutoEnzyme{<:Union{ForwardMode,Nothing}}, x)
+    B = pick_batchsize(backend, length(x))
+    shadow = chunkedonehot(x, Val(B))
+    return EnzymeForwardOneArgJacobianExtras{B,typeof(shadow)}(shadow)
 end
 
 function DI.jacobian(
     f,
     backend::AutoEnzyme{<:Union{ForwardMode,Nothing}},
     x,
-    extras::EnzymeForwardOneArgJacobianExtras{C},
-) where {C}
-    jac_wrongshape = jacobian(forward_mode(backend), f, x, Val{C}(); shadow=extras.shadow)
+    extras::EnzymeForwardOneArgJacobianExtras{B},
+) where {B}
+    jac_wrongshape = jacobian(forward_mode(backend), f, x, Val(B); shadow=extras.shadow)
     nx = length(x)
     ny = length(jac_wrongshape) ÷ length(x)
     return reshape(jac_wrongshape, ny, nx)

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/reverse_onearg.jl

@@ -160,22 +160,22 @@ end
 
 #=
 
-struct EnzymeReverseOneArgJacobianExtras{C,N} end
+struct EnzymeReverseOneArgJacobianExtras{B,N} end
 
-function DI.prepare_jacobian(f, ::AutoReverseEnzyme, x)
-    C = pick_chunksize(length(x))
+function DI.prepare_jacobian(f, backend::AutoReverseEnzyme, x)
+    B = pick_batchsize(backend, length(x))
     y = f(x)
     N = length(y)
-    return EnzymeReverseOneArgJacobianExtras{C,N}()
+    return EnzymeReverseOneArgJacobianExtras{B,N}()
 end
 
 function DI.jacobian(
     f,
     backend::AutoReverseEnzyme,
     x::AbstractArray,
     ::EnzymeReverseOneArgJacobianExtras{C,N},
-) where {C,N}
-    jac_wrongshape = jacobian(reverse_mode(backend), f, x, Val{N}(), Val{C}())
+) where {B,N}
+    jac_wrongshape = jacobian(reverse_mode(backend), f, x, Val(N), Val(B))
     nx = length(x)
     ny = length(jac_wrongshape) ÷ length(x)
     jac_rightshape = reshape(jac_wrongshape, ny, nx)

DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/DifferentiationInterfaceForwardDiffExt.jl

@@ -38,6 +38,14 @@ using LinearAlgebra: dot, mul!
 
 DI.check_available(::AutoForwardDiff) = true
 
+function DI.pick_batchsize(::AutoForwardDiff{C}, dimension::Integer) where {C}
+    if isnothing(C)
+        return ForwardDiff.pickchunksize(dimension)
+    else
+        return min(dimension, C)
+    end
+end
+
 include("utils.jl")
 include("onearg.jl")
 include("twoarg.jl")

DifferentiationInterface/ext/DifferentiationInterfaceZygoteExt/DifferentiationInterfaceZygoteExt.jl

@@ -3,6 +3,7 @@ module DifferentiationInterfaceZygoteExt
 using ADTypes: AutoForwardDiff, AutoZygote
 import DifferentiationInterface as DI
 using DifferentiationInterface:
+    Batch,
     HVPExtras,
     NoGradientExtras,
     NoHessianExtras,
@@ -103,20 +104,47 @@ struct ZygoteHVPExtras{G,PE} <: HVPExtras
     pushforward_extras::PE
 end
 
-function DI.prepare_hvp(f, ::AutoZygote, x, v)
+function DI.prepare_hvp(f, ::AutoZygote, x, dx)
     ∇f(x) = only(gradient(f, x))
-    pushforward_extras = DI.prepare_pushforward(∇f, AutoForwardDiff(), x, v)
+    pushforward_extras = DI.prepare_pushforward(∇f, AutoForwardDiff(), x, dx)
     return ZygoteHVPExtras(∇f, pushforward_extras)
 end
 
-function DI.hvp(f, ::AutoZygote, x, v, extras::ZygoteHVPExtras)
+function DI.hvp(f, ::AutoZygote, x, dx, extras::ZygoteHVPExtras)
     @compat (; ∇f, pushforward_extras) = extras
-    return DI.pushforward(∇f, AutoForwardDiff(), x, v, pushforward_extras)
+    return DI.pushforward(∇f, AutoForwardDiff(), x, dx, pushforward_extras)
 end
 
-function DI.hvp!(f, p, ::AutoZygote, x, v, extras::ZygoteHVPExtras)
+function DI.hvp!(f, dg, ::AutoZygote, x, dx, extras::ZygoteHVPExtras)
     @compat (; ∇f, pushforward_extras) = extras
-    return DI.pushforward!(∇f, p, AutoForwardDiff(), x, v, pushforward_extras)
+    return DI.pushforward!(∇f, dg, AutoForwardDiff(), x, dx, pushforward_extras)
+end
+
+struct ZygoteHVPBatchedExtras{G,PE} <: HVPExtras
+    ∇f::G
+    pushforward_batched_extras::PE
+end
+
+function DI.prepare_hvp_batched(f, ::AutoZygote, x, dx::Batch)
+    ∇f(x) = only(gradient(f, x))
+    pushforward_batched_extras = DI.prepare_pushforward_batched(
+        ∇f, AutoForwardDiff(), x, dx
+    )
+    return ZygoteHVPBatchedExtras(∇f, pushforward_batched_extras)
+end
+
+function DI.hvp_batched(f, ::AutoZygote, x, dx::Batch, extras::ZygoteHVPBatchedExtras)
+    @compat (; ∇f, pushforward_batched_extras) = extras
+    return DI.pushforward_batched(∇f, AutoForwardDiff(), x, dx, pushforward_batched_extras)
+end
+
+function DI.hvp_batched!(
+    f, dg::Batch, ::AutoZygote, x, dx::Batch, extras::ZygoteHVPBatchedExtras
+)
+    @compat (; ∇f, pushforward_batched_extras) = extras
+    return DI.pushforward_batched!(
+        ∇f, dg, AutoForwardDiff(), x, dx, pushforward_batched_extras
+    )
 end
 
 ## Hessian

DifferentiationInterface/src/DifferentiationInterface.jl

@@ -50,8 +50,7 @@ include("second_order/second_order.jl")
 
 include("utils/traits.jl")
 include("utils/basis.jl")
-include("utils/printing.jl")
-include("utils/chunk.jl")
+include("utils/batch.jl")
 include("utils/check.jl")
 include("utils/exceptions.jl")
 include("utils/maybe.jl")
@@ -73,6 +72,9 @@ include("sparse/hessian.jl")
 
 include("misc/differentiate_with.jl")
 include("misc/sparsity_detector.jl")
+include("misc/from_primitive.jl")
+
+include("utils/printing.jl")
 
 function __init__()
     @require_extensions