separated VJP implementation for AutoZygote, FiniteDiff

Author: vpuri3

ext/SparseDiffToolsZygote.jl

@@ -1,16 +1,12 @@
 module SparseDiffToolsZygote
 
-if isdefined(Base, :get_extension)
-    import Zygote
-    using LinearAlgebra
-    using SparseDiffTools: SparseDiffTools, DeivVecTag
-    using ForwardDiff: ForwardDiff, Dual, partials
-else
-    import ..Zygote
-    using ..LinearAlgebra
-    using ..SparseDiffTools: SparseDiffTools, DeivVecTag
-    using ..ForwardDiff: ForwardDiff, Dual, partials
-end
+import Zygote
+using ADTypes
+using LinearAlgebra
+using SparseDiffTools: SparseDiffTools, DeivVecTag, AutoDiffVJP
+using ForwardDiff: ForwardDiff, Dual, partials
+import SciMLOperators: update_coefficients, update_coefficients!
+import Setfield: @set!
 
 ### Jac, Hes products
 
@@ -75,14 +71,56 @@ end
 
 ## VecJac products
 
+# VJP methods
 function SparseDiffTools.auto_vecjac!(du, f, x, v)
     !hasmethod(f, (typeof(x),)) && error("For inplace function use autodiff = AutoFiniteDiff()")
     du .= reshape(SparseDiffTools.auto_vecjac(f, x, v), size(du))
 end
 
 function SparseDiffTools.auto_vecjac(f, x, v)
-    vv, back = Zygote.pullback(f, x)
-    return vec(back(reshape(v, size(vv)))[1])
+    y, back = Zygote.pullback(f, x)
+    return vec(back(reshape(v, size(y)))[1])
+end
+
+# overload operator interface
+function SparseDiffTools._vecjac(f, u, autodiff::AutoZygote)
+
+    cache = ()
+    pullback = Zygote.pullback(f, u)
+
+    AutoDiffVJP(f, u, cache, autodiff, pullback)
+end
+
+function update_coefficients(L::AutoDiffVJP{AD}, u, p, t) where{AD <: AutoZygote}
+    @set! L.f = update_coefficients(L.f, u, p, t)
+    @set! L.u = u
+    @set! L.pullback = Zygote.pullback(L.f, u)
+end
+
+function update_coefficients!(L::AutoDiffVJP{AD}, u, p, t) where{AD <: AutoZygote}
+    update_coefficients!(L.f, u, p, t)
+    copy!(L.u, u)
+    L.pullback = Zygote.pullback(L.f, u)
+    L
+end
+
+# Interpret the call as df/du' * v
+function (L::AutoDiffVJP{AD})(v, p, t) where{AD <: AutoZygote}
+
+    y, back = L.pullback
+    V = reshape(v, size(y))
+
+    back(V)[1] |> vec
+end
+
+# prefer non in-place method
+function (L::AutoDiffVJP{AD, IIP, true})(dv, v, p, t) where {AD <: AutoZygote, IIP}
+    _dv = L(v, p, t)
+    copy!(dv, _dv)
+end
+
+function (L::AutoDiffVJP{AD, true, false})(dv, v, p, t) where {AD <: AutoZygote}
+    SparseDiffTools.auto_vecjac!(dv, L.f, L.u, v, L.cache...)
 end
 
 end # module

src/differentiation/vecjac_products.jl

@@ -37,15 +37,35 @@ end
 
 ### Operator Forms
 
-struct RevModeAutoDiffVecProd{ad, iip, oop, F, U, C, V, V!} <: AbstractAutoDiffVecProd
+"""
+    VecJac(f, u, [p, t]; autodiff = AutoFiniteDiff())
+"""
+function VecJac(f, u::AbstractArray, p = nothing, t = nothing;
+                autodiff = AutoFiniteDiff(), kwargs...)
+
+    L = _vecjac(f, u, autodiff)
+    IIP, OOP = get_iip_oop(L)
+
+    FunctionOperator(L, u, u; isinplace = IIP, outofplace = OOP,
+                     p = p, t = t, islinear = true, kwargs...)
+end
+
+function _vecjac(f, u, autodiff::AutoFiniteDiff)
+
+    cache = (similar(u), similar(u))
+    pullback = nothing
+
+    AutoDiffVJP(f, u, cache, autodiff, pullback)
+end
+
+mutable struct AutoDiffVJP{AD, IIP, OOP, F, U, C, PB} <: AbstractAutoDiffVecProd
     f::F
     u::U
     cache::C
-    vecprod::V
-    vecprod!::V!
-    autodiff::ad
+    autodiff::AD
+    pullback::PB
 
-    function RevModeAutoDiffVecProd(f, u, cache, vecprod, vecprod!, autodiff)
+    function AutoDiffVJP(f, u, cache, autodiff, pullback)
 
         outofplace = static_hasmethod(f, typeof((u,)))
         isinplace = static_hasmethod(f, typeof((u, u)))
@@ -62,83 +82,45 @@ struct RevModeAutoDiffVecProd{ad, iip, oop, F, U, C, V, V!} <: AbstractAutoDiffV
             typeof(f),
             typeof(u),
             typeof(cache),
-            typeof(vecprod),
-            typeof(vecprod!)
+            typeof(pullback),
            }(
-             f, u, cache, vecprod, vecprod!, autodiff,
+             f, u, cache, autodiff, pullback,
             )
     end
 end
 
-function get_iip_oop(::RevModeAutoDiffVecProd{ad, iip, oop}) where{ad, iip, oop}
-    iip, oop
+function get_iip_oop(::AutoDiffVJP{AD, IIP, OOP}) where{AD, IIP, OOP}
+    IIP, OOP
 end
 
-function update_coefficients(L::RevModeAutoDiffVecProd, u, p, t)
+function update_coefficients(L::AutoDiffVJP{AD}, u, p, t) where{AD <: AutoFiniteDiff}
     @set! L.f = update_coefficients(L.f, u, p, t)
     @set! L.u = u
 end
 
-function update_coefficients!(L::RevModeAutoDiffVecProd, u, p, t)
+function update_coefficients!(L::AutoDiffVJP{AD}, u, p, t) where{AD <: AutoFiniteDiff}
     update_coefficients!(L.f, u, p, t)
     copy!(L.u, u)
     L
 end
 
-# Interpret the call as df/du' * u
-function (L::RevModeAutoDiffVecProd)(v, p, t)
-    L.vecprod(L.f, L.u, v)
+# Interpret the call as df/du' * v
+function (L::AutoDiffVJP{AD})(v, p, t) where{AD <: AutoFiniteDiff}
+    num_vecjac(L.f, L.u, v)
 end
 
-# prefer non in-place method
-function (L::RevModeAutoDiffVecProd{ad, iip, true})(dv, v, p, t) where {ad, iip}
-    L.vecprod!(dv, L.f, L.u, v, L.cache...)
+function (L::AutoDiffVJP{AD})(dv, v, p, t) where{AD <: AutoFiniteDiff}
+    num_vecjac!(dv, L.f, L.u, v, L.cache...)
 end
 
-function (L::RevModeAutoDiffVecProd{ad, true, false})(dv, v, p, t) where {ad}
-    L.vecprod!(dv, L.f, L.u, v, L.cache...)
-end
-
-function Base.resize!(L::RevModeAutoDiffVecProd, n::Integer)
+function Base.resize!(L::AutoDiffVJP, n::Integer)
 
     static_hasmethod(resize!, typeof((L.f, n))) && resize!(L.f, n)
     resize!(L.u, n)
 
     for v in L.cache
         resize!(v, n)
     end
-end
-
-"""
-    VecJac(f, u, [p, t]; autodiff = AutoFiniteDiff())
-
-Returns FunctionOperator that computes
-"""
-function VecJac(f, u::AbstractArray, p = nothing, t = nothing;
-                autodiff = AutoFiniteDiff(), kwargs...)
-
-    vecprod, vecprod!, cache = if autodiff isa AutoFiniteDiff
-        num_vecjac, num_vecjac!, (similar(u), similar(u))
-    elseif autodiff isa AutoZygote
-        @assert static_hasmethod(auto_vecjac, typeof((f, u, u))) "To use AutoZygote() AD, first load Zygote with `using Zygote`, or `import Zygote`"
-
-        auto_vecjac, auto_vecjac!, ()
-    end
-
-    L = RevModeAutoDiffVecProd(f, u, cache, vecprod, vecprod!, autodiff)
-
-    iip, oop = get_iip_oop(L)
-
-    FunctionOperator(L, u, u; isinplace = iip, outofplace = oop,
-                     p = p, t = t, islinear = true, kwargs...)
-end
-
-
-function FixedVecJac(f, u::AbstractArray, p = nothing, t = nothing;
-                     autodiff = AutoFiniteDiff(), kwargs...)
-    _fixedvecjac(f, u, p, t, autodiff, kwargs)
-end
 
-function _fixedvecjac(f, u, p, t, ad::AutoFiniteDiff, kwargs)
 end
 #