Merge pull request #245 from vpuri3/vj

Fix VecJac

Author: ChrisRackauckas

Project.toml

@@ -16,12 +16,19 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961"
+Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 StaticArrayInterface = "0d7ed370-da01-4f52-bd93-41d350b8b718"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Tricks = "410a4b4d-49e4-4fbc-ab6d-cb71b17b3775"
 VertexSafeGraphs = "19fa3120-7c27-5ec5-8db8-b0b0aa330d6f"
 
+[weakdeps]
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
+
+[extensions]
+SparseDiffToolsZygote = "Zygote"
+
 [compat]
 ADTypes = "0.1"
 Adapt = "3.0"
@@ -33,17 +40,14 @@ ForwardDiff = "0.10"
 Graphs = "1"
 Reexport = "1"
 Requires = "1"
-SciMLOperators = "0.1.19, 0.2"
+SciMLOperators = "0.2.10"
 StaticArrayInterface = "1.3"
 StaticArrays = "1"
 Tricks = "0.1.6"
 VertexSafeGraphs = "0.2"
 Zygote = "0.6"
 julia = "1.6"
 
-[extensions]
-SparseDiffToolsZygote = "Zygote"
-
 [extras]
 ArrayInterfaceBandedMatrices = "2e50d22c-5be1-4042-81b1-c572ed69783d"
 ArrayInterfaceBlockBandedMatrices = "5331f1e9-51c7-46b0-a9b0-df4434785e0a"
@@ -60,6 +64,3 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
 test = ["Test", "ArrayInterfaceBandedMatrices", "ArrayInterfaceBlockBandedMatrices", "BandedMatrices", "BlockBandedMatrices", "IterativeSolvers", "Pkg", "Random", "SafeTestsets", "Symbolics", "Zygote", "StaticArrays"]
-
-[weakdeps]
-Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"

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,70 @@ end
 
 ## VecJac products
 
-function SparseDiffTools.auto_vecjac!(du, f, x, v, cache1 = nothing, cache2 = nothing)
-    !hasmethod(f, (typeof(x),)) && error("For inplace function use autodiff = false")
+# 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; VJP_input = nothing,
+                            ) where{AD <: AutoZygote}
+
+    if !isnothing(VJP_input)
+        @set! L.u = VJP_input
+    end
+
+    @set! L.f = update_coefficients(L.f, L.u, p, t)
+    @set! L.pullback = Zygote.pullback(L.f, L.u)
+end
+
+function update_coefficients!(L::AutoDiffVJP{AD}, u, p, t; VJP_input = nothing,
+                             ) where{AD <: AutoZygote}
+
+    if !isnothing(VJP_input)
+        copy!(L.u, VJP_input)
+    end
+
+    update_coefficients!(L.f, L.u, p, t)
+    L.pullback = Zygote.pullback(L.f, L.u)
+
+    L
+end
+
+# Interpret the call as df/du' * v
+function (L::AutoDiffVJP{AD})(v, p, t; VJP_input = nothing) where{AD <: AutoZygote}
+    # ignore VJP_input as pullback was computed in update_coefficients(...)
+
+    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; VJP_input = nothing) where {AD <: AutoZygote, IIP}
+    # ignore VJP_input as pullback was computed in update_coefficients!(...)
+
+    _dv = L(v, p, t; VJP_input = VJP_input)
+    copy!(dv, _dv)
+end
+
+function (L::AutoDiffVJP{AD, true, false})(dv, v, p, t; VJP_input = nothing) where {AD <: AutoZygote}
+    @error("Zygote requires an out of place method with signature f(u).")
 end
 
 end # module

src/SparseDiffTools.jl

@@ -7,6 +7,7 @@ using Graphs
 using Graphs: SimpleGraph
 using VertexSafeGraphs
 using Adapt
+
 using Reexport
 @reexport using ADTypes
 
@@ -23,6 +24,7 @@ using ArrayInterface: matrix_colors
 using SciMLOperators
 import SciMLOperators: update_coefficients, update_coefficients!
 using Tricks: Tricks, static_hasmethod
+using Setfield: @set!
 
 abstract type AbstractAutoDiffVecProd end
 

src/differentiation/jaches_products.jl

@@ -210,8 +210,8 @@ struct FwdModeAutoDiffVecProd{F, U, C, V, V!} <: AbstractAutoDiffVecProd
 end
 
 function update_coefficients(L::FwdModeAutoDiffVecProd, u, p, t)
-    f = update_coefficients(L.f, u, p, t)
-    FwdModeAutoDiffVecProd(f, u, L.cache, L.vecprod, L.vecprod!)
+    @set! L.f = update_coefficients(L.f, u, p, t)
+    @set! L.u = u
 end
 
 function update_coefficients!(L::FwdModeAutoDiffVecProd, u, p, t)
@@ -248,7 +248,7 @@ function JacVec(f, u::AbstractArray, p = nothing, t = nothing;
     elseif autodiff isa AutoForwardDiff
         cache1 = Dual{
                       typeof(ForwardDiff.Tag(tag, eltype(u))), eltype(u), 1
-                      }.(u, ForwardDiff.Partials.(tuple.(u)))
+                     }.(u, ForwardDiff.Partials.(tuple.(u)))
 
         cache2 = copy(cache1)
 

src/differentiation/vecjac_products.jl

@@ -37,17 +37,67 @@ end
 
 ### Operator Forms
 
-struct RevModeAutoDiffVecProd{ad, iip, oop, F, U, C, V, V!} <: AbstractAutoDiffVecProd
+"""
+    VecJac(f, u, [p, t]; autodiff = AutoFiniteDiff())
+
+Returns SciMLOperators.FunctionOperator which computes vector-jacobian
+product `df/du * v`.
+
+```
+L = VecJac(f, u)
+
+L * v         # = df/du * v
+mul!(w, L, v) # = df/du * v
+
+L(v, p, t; VJP_input = w)    # = df/dw * v
+L(x, v, p, t; VJP_input = w) # = df/dw * v
+```
+"""
+function VecJac(f, u::AbstractArray, p = nothing, t = nothing;
+                autodiff = AutoFiniteDiff(), kwargs...)
+
+    L = _vecjac(f, u, autodiff)
+    IIP, OOP = get_iip_oop(L)
+
+    if isa(autodiff, AutoZygote) & !OOP
+        msg = "Zygote requires an out of place method with signature f(u)."
+        throw(ArgumentError(msg))
+    end
+
+    FunctionOperator(L, u, u; isinplace = IIP, outofplace = OOP,
+                     p = p, t = t, islinear = true,
+                     accepted_kwargs = (:VJP_input,), 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
+    """ Compute VJP of `f` at `u`, applied to vector `v`: `df/du' * u` """
     f::F
+    """ input to `f` """
     u::U
+    """ Cache for num_vecjac! when autodiff isa AutoFintieDiff """
     cache::C
-    vecprod::V
-    vecprod!::V!
+    """ Type of automatic differentiation algorithm """
+    autodiff::AD
+    """ stores the result of Zygote.pullback for AutoZygote """
+    pullback::PB
+
+    function AutoDiffVJP(f, u, cache, autodiff, pullback)
 
-    function RevModeAutoDiffVecProd(f, u, cache, vecprod, vecprod!;
-                                    autodiff = AutoFiniteDiff(),
-                                    isinplace = false, outofplace = true)
-        @assert isinplace || outofplace
+        outofplace = static_hasmethod(f, typeof((u,)))
+        isinplace = static_hasmethod(f, typeof((u, u)))
+
+        if !(isinplace) & !(outofplace)
+            msg = "$f must have signature f(u), or f(du, u)"
+            throw(ArgumentError(msg))
+        end
 
         new{
             typeof(autodiff),
@@ -56,72 +106,58 @@ struct RevModeAutoDiffVecProd{ad, iip, oop, F, U, C, V, V!} <: AbstractAutoDiffV
             typeof(f),
             typeof(u),
             typeof(cache),
-            typeof(vecprod),
-            typeof(vecprod!)
-            }(f, u, cache, vecprod, vecprod!)
+            typeof(pullback),
+           }(
+             f, u, cache, autodiff, pullback,
+            )
     end
 end
 
-function update_coefficients(L::RevModeAutoDiffVecProd, u, p, t)
-    f = update_coefficients(L.f, u, p, t)
-    RevModeAutoDiffVecProd(f, u, L.vecprod, L.vecprod!, L.cache)
+function get_iip_oop(::AutoDiffVJP{AD, IIP, OOP}) where{AD, IIP, OOP}
+    IIP, OOP
 end
 
-function update_coefficients!(L::RevModeAutoDiffVecProd, u, p, t)
-    update_coefficients!(L.f, u, p, t)
-    copy!(L.u, u)
-    L
+function update_coefficients(L::AutoDiffVJP{AD}, u, p, t; VJP_input = nothing,
+                            ) where{AD <: AutoFiniteDiff}
+
+    if !isnothing(VJP_input)
+        @set! L.u = VJP_input
+    end
+
+    @set! L.f = update_coefficients(L.f, L.u, p, t)
 end
 
-# Interpret the call as df/du' * u
-function (L::RevModeAutoDiffVecProd)(v, p, t)
-    L.vecprod(L.f, L.u, v)
+function update_coefficients!(L::AutoDiffVJP{AD}, u, p, t; VJP_input = nothing,
+                             ) where{AD <: AutoFiniteDiff}
+
+    if !isnothing(VJP_input)
+        copy!(L.u, VJP_input)
+    end
+
+    update_coefficients!(L.f, L.u, p, t)
+
+    L
 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...)
+# Interpret the call as df/du' * v
+function (L::AutoDiffVJP{AD})(v, p, t; VJP_input = nothing,) where{AD <: AutoFiniteDiff}
+    # ignore VJP_input as L.u was set in update_coefficients(...)
+    num_vecjac(L.f, L.u, v)
 end
 
-function (L::RevModeAutoDiffVecProd{ad, true, false})(dv, v, p, t) where {ad}
-    L.vecprod!(dv, L.f, L.u, v, L.cache...)
+function (L::AutoDiffVJP{AD})(dv, v, p, t; VJP_input = nothing,) where{AD <: AutoFiniteDiff}
+    # ignore VJP_input as L.u was set in update_coefficients!(...)
+    num_vecjac!(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
-
-function VecJac(f, u::AbstractArray, p = nothing, t = nothing; autodiff = AutoFiniteDiff(),
-                kwargs...)
-    vecprod, vecprod! = if autodiff isa AutoFiniteDiff
-        num_vecjac, num_vecjac!
-    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
-
-    cache = (similar(u), similar(u))
-
-    outofplace = static_hasmethod(f, typeof((u,)))
-    isinplace = static_hasmethod(f, typeof((u, u)))
-
-    if !(isinplace) & !(outofplace)
-        error("$f must have signature f(u), or f(du, u)")
-    end
-
-    L = RevModeAutoDiffVecProd(f, u, cache, vecprod, vecprod!; autodiff = autodiff,
-                               isinplace = isinplace, outofplace = outofplace)
-
-    FunctionOperator(L, u, u;
-                     isinplace = isinplace, outofplace = outofplace,
-                     p = p, t = t, islinear = true,
-                     kwargs...)
 end
 #

test/test_jaches_products.jl

@@ -143,7 +143,7 @@ L = JacVec(f, copy(x), 1.0, 1.0; tag = MyTag())
 
 # Resize test
 for M in (100, 400)
-    L = JacVec(f2, copy(x), 1.0, 1.0)
+    local L = JacVec(f2, copy(x), 1.0, 1.0)
     resize!(L, M)
     _x = resize!(copy(x), M)
     _u = rand(M)