fix VecJac constructor, clean auto_vecjac<bang>, add setfield

Author: vpuri3

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"
@@ -41,9 +48,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"

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

@@ -43,11 +43,17 @@ struct RevModeAutoDiffVecProd{ad, iip, oop, F, U, C, V, V!} <: AbstractAutoDiffV
     cache::C
     vecprod::V
     vecprod!::V!
+    autodiff::ad
 
-    function RevModeAutoDiffVecProd(f, u, cache, vecprod, vecprod!;
-                                    autodiff = AutoFiniteDiff(),
-                                    isinplace = false, outofplace = true)
-        @assert isinplace || outofplace
+    function RevModeAutoDiffVecProd(f, u, cache, vecprod, vecprod!, autodiff)
+
+        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),
@@ -58,13 +64,19 @@ struct RevModeAutoDiffVecProd{ad, iip, oop, F, U, C, V, V!} <: AbstractAutoDiffV
             typeof(cache),
             typeof(vecprod),
             typeof(vecprod!)
-            }(f, u, cache, vecprod, vecprod!)
+           }(
+             f, u, cache, vecprod, vecprod!, autodiff,
+            )
     end
 end
 
+function get_iip_oop(::RevModeAutoDiffVecProd{ad, iip, oop}) where{ad, iip, oop}
+    iip, oop
+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)
+    @set! L.f = update_coefficients(L.f, u, p, t)
+    @set! L.u = u
 end
 
 function update_coefficients!(L::RevModeAutoDiffVecProd, u, p, t)
@@ -97,31 +109,36 @@ function Base.resize!(L::RevModeAutoDiffVecProd, n::Integer)
     end
 end
 
-function VecJac(f, u::AbstractArray, p = nothing, t = nothing; autodiff = AutoFiniteDiff(),
-                kwargs...)
-    vecprod, vecprod! = if autodiff isa AutoFiniteDiff
-        num_vecjac, num_vecjac!
+"""
+    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!
+        auto_vecjac, auto_vecjac!, ()
     end
 
-    cache = (similar(u), similar(u))
+    L = RevModeAutoDiffVecProd(f, u, cache, vecprod, vecprod!, autodiff)
 
-    outofplace = static_hasmethod(f, typeof((u,)))
-    isinplace = static_hasmethod(f, typeof((u, u)))
+    iip, oop = get_iip_oop(L)
 
-    if !(isinplace) & !(outofplace)
-        error("$f must have signature f(u), or f(du, u)")
-    end
+    FunctionOperator(L, u, u; isinplace = iip, outofplace = oop,
+                     p = p, t = t, islinear = true, kwargs...)
+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...)
+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
 #