rm ZygoteVecJac. rename ZygoteHesVec to BackHesVec, and move it to main

Author: vpuri3

ext/SparseDiffToolsZygote.jl

@@ -3,19 +3,13 @@ module SparseDiffToolsZygote
 if isdefined(Base, :get_extension)
     import Zygote
     using LinearAlgebra
-    using SparseDiffTools: SparseDiffTools, DeivVecTag, FwdModeAutoDiffVecProd, VecJac
+    using SparseDiffTools: SparseDiffTools, DeivVecTag
     using ForwardDiff: ForwardDiff, Dual, partials
-    using SciMLOperators: FunctionOperator
-    using Tricks: static_hasmethod
-    using ADTypes
 else
     import ..Zygote
     using ..LinearAlgebra
-    using ..SparseDiffTools: SparseDiffTools, DeivVecTag, FwdModeAutoDiffVecProd, VecJac
+    using ..SparseDiffTools: SparseDiffTools, DeivVecTag
     using ..ForwardDiff: ForwardDiff, Dual, partials
-    using ..SciMLOperators: FunctionOperator
-    using ..Tricks: static_hasmethod
-    using ..ADTypes
 end
 
 ### Jac, Hes products
@@ -71,39 +65,6 @@ function SparseDiffTools.autoback_hesvec(f, x, v)
     ForwardDiff.partials.(g(y), 1)
 end
 
-# Operator Forms
-
-function SparseDiffTools.ZygoteHesVec(f, u::AbstractArray, p = nothing, t = nothing; autodiff = AutoZygote())
-
-    cache, vecprod, vecprod! = if autodiff isa AutoFiniteDiff
-        cache1 = similar(u)
-        cache2 = similar(u)
-
-        (cache1, cache2), SparseDiffTools.numback_hesvec, SparseDiffTools.numback_hesvec!
-    elseif autodiff isa AutoZygote
-        cache1 = Dual{
-                      typeof(ForwardDiff.Tag(DeivVecTag(),eltype(u))), eltype(u), 1
-                     }.(u, ForwardDiff.Partials.(tuple.(u)))
-        cache2 = copy(u)
-
-        (cache1, cache2), SparseDiffTools.autoback_hesvec, SparseDiffTools.autoback_hesvec!
-    end
-
-    outofplace = static_hasmethod(f, typeof((u,)))
-    isinplace  = static_hasmethod(f, typeof((u,)))
-
-    if !(isinplace) & !(outofplace)
-        error("$f must have signature f(u).")
-    end
-
-    L = FwdModeAutoDiffVecProd(f, u, cache, vecprod, vecprod!)
-
-    FunctionOperator(L, u, u;
-                     isinplace = isinplace, outofplace = outofplace,
-                     p = p, t = t, islinear = true,
-                    )
-end
-
 ## VecJac products
 
 function SparseDiffTools.auto_vecjac!(du, f, x, v, cache1 = nothing, cache2 = nothing)
@@ -116,8 +77,4 @@ function SparseDiffTools.auto_vecjac(f, x, v)
     return vec(back(reshape(v, size(vv)))[1])
 end
 
-function SparseDiffTools.ZygoteVecJac(args...; kwargs...)
-    VecJac(args...; autodiff = AutoZygote(), kwargs...)
-end
-
 end # module

src/SparseDiffTools.jl

@@ -49,7 +49,7 @@ export contract_color,
        autonum_hesvec, autonum_hesvec!,
        num_hesvecgrad, num_hesvecgrad!,
        auto_hesvecgrad, auto_hesvecgrad!,
-       JacVec, HesVec, HesVecGrad, VecJac,
+       JacVec, HesVec, BackHesVec, HesVecGrad, VecJac,
        update_coefficients, update_coefficients!,
        value!
 
@@ -78,8 +78,6 @@ function autoback_hesvec end
 function autoback_hesvec! end
 function auto_vecjac end
 function auto_vecjac! end
-function ZygoteVecJac end
-function ZygoteHesVec end
 
 @static if !isdefined(Base, :get_extension)
     function __init__()
@@ -93,7 +91,6 @@ end
 export
        numback_hesvec, numback_hesvec!,
        autoback_hesvec, autoback_hesvec!,
-       auto_vecjac, auto_vecjac!,
-       ZygoteVecJac, ZygoteHesVec
+       auto_vecjac, auto_vecjac!
 
 end # module

src/differentiation/jaches_products.jl

@@ -290,6 +290,39 @@ function HesVec(f, u::AbstractArray, p = nothing, t = nothing; autodiff = AutoFo
                     )
 end
 
+function BackHesVec(f, u::AbstractArray, p = nothing, t = nothing; autodiff = AutoFiniteDiff())
+
+    cache, vecprod, vecprod! = if autodiff isa AutoFiniteDiff
+        cache1 = similar(u)
+        cache2 = similar(u)
+
+        (cache1, cache2), numback_hesvec, numback_hesvec!
+    elseif autodiff isa AutoZygote
+        @assert static_hasmethod(autoback_hesvec, typeof((f, u, u))) "To use AutoZygote() AD, first load Zygote with `using Zygote`, or `import Zygote`"
+
+        cache1 = Dual{
+                      typeof(ForwardDiff.Tag(DeivVecTag(),eltype(u))), eltype(u), 1
+                     }.(u, ForwardDiff.Partials.(tuple.(u)))
+        cache2 = copy(u)
+
+        (cache1, cache2), autoback_hesvec, autoback_hesvec!
+    end
+
+    outofplace = static_hasmethod(f, typeof((u,)))
+    isinplace  = static_hasmethod(f, typeof((u,)))
+
+    if !(isinplace) & !(outofplace)
+        error("$f must have signature f(u).")
+    end
+
+    L = FwdModeAutoDiffVecProd(f, u, cache, vecprod, vecprod!)
+
+    FunctionOperator(L, u, u;
+                     isinplace = isinplace, outofplace = outofplace,
+                     p = p, t = t, islinear = true,
+                    )
+end
+
 function HesVecGrad(f, u::AbstractArray, p = nothing, t = nothing; autodiff = AutoForwardDiff())
 
     cache, vecprod, vecprod! = if autodiff isa AutoFiniteDiff

test/test_jaches_products.jl

@@ -113,12 +113,12 @@ dy=rand(N);_dy=copy(dy);@test mul!(dy,L,v,a,b)≈a*numauto_hesvec(g,x,v)+b*_dy r
 out = similar(v)
 gmres!(out, L, v)
 
-@info "ZygoteHesVec"
+@info "BackHesVec"
 using Zygote
 x = rand(N)
 v = rand(N)
 
-L = ZygoteHesVec(g, x, autodiff = AutoFiniteDiff())
+L = BackHesVec(g, x)
 @test L * x ≈ numback_hesvec(g, x, x) rtol = 1e-2
 @test L * v ≈ numback_hesvec(g, x, v) rtol = 1e-2
 @test mul!(dy, L, v)≈numback_hesvec(g, x, v) rtol=1e-2
@@ -127,7 +127,7 @@ update_coefficients!(L, v, nothing, 0.0)
 @test mul!(dy, L, v)≈numback_hesvec(g, v, v) rtol=1e-2
 dy=rand(N);_dy=copy(dy);@test mul!(dy,L,v,a,b) ≈ a*numback_hesvec(g,x,v) + b*_dy rtol=1e-2
 
-L = ZygoteHesVec(g, x)
+L = BackHesVec(g, x, autodiff = AutoZygote())
 @test L * x ≈ autoback_hesvec(g, x, x)
 @test L * v ≈ autoback_hesvec(g, x, v)
 @test mul!(dy, L, v)≈autoback_hesvec(g, x, v) rtol=1e-8