zygote ext file

Author: vpuri3

…ifferentiation/jaches_products_zygote.jl ext/SparseDiffToolsZygote.jlsrc/differentiation/jaches_products_zygote.jl renamed to ext/SparseDiffToolsZygote.jl src/differentiation/jaches_products_zygote.jl renamed to ext/SparseDiffToolsZygote.jl

@@ -1,3 +1,24 @@
+module SparseDiffToolsZygote
+
+import Zygote
+
+using SparseDiffTools
+using SparseDiffTools: DeviVecTag, FwdModeAutoDiffVecProd
+
+using ForwardDiff
+using ForwardDiff: Dual, Tag
+
+using SciMLOperators: FunctionOperator
+using Tricks: static_hasmethod
+
+export
+       numback_hesvec, numback_hesvec!,
+       autoback_hesvec, autoback_hesvec!,
+       auto_vecjac, auto_vecjac!,
+       ZygoteVecJac, ZygoteHesVec
+
+### Jac, Hes products
+
 function numback_hesvec!(dy, f, x, v, cache1 = similar(v), cache2 = similar(v))
     g = let f = f
         (dx, x) -> dx .= first(Zygote.gradient(f, x))
@@ -49,7 +70,7 @@ function autoback_hesvec(f, x, v)
     ForwardDiff.partials.(g(y), 1)
 end
 
-### Operator Forms
+# Operator Forms
 
 function ZygoteHesVec(f, u::AbstractArray, p = nothing, t = nothing; autodiff = true)
 
@@ -82,4 +103,19 @@ function ZygoteHesVec(f, u::AbstractArray, p = nothing, t = nothing; autodiff =
                      p = p, t = t, islinear = true,
                     )
 end
-#
+
+## VecJac products
+
+function auto_vecjac!(du, f, x, v, cache1 = nothing, cache2 = nothing)
+    !hasmethod(f, (typeof(x),)) && error("For inplace function use autodiff = false")
+    du .= reshape(auto_vecjac(f, x, v), size(du))
+end
+
+function auto_vecjac(f, x, v)
+    vv, back = Zygote.pullback(f, x)
+    return vec(back(reshape(v, size(vv)))[1])
+end
+
+ZygoteVecJac(args...; autodiff = true, kwargs...) = VecJac(args...; autodiff = autodiff, kwargs...)
+
+end # module

src/SparseDiffTools.jl

@@ -5,7 +5,6 @@ using FiniteDiff
 using ForwardDiff
 using Graphs
 using Graphs: SimpleGraph
-using Requires
 using VertexSafeGraphs
 using Adapt
 
@@ -69,16 +68,16 @@ Base.@pure __parameterless_type(T) = Base.typename(T).wrapper
 parameterless_type(x) = parameterless_type(typeof(x))
 parameterless_type(x::Type) = __parameterless_type(x)
 
-function __init__()
-    @require Zygote="e88e6eb3-aa80-5325-afca-941959d7151f" begin
-        export numback_hesvec, numback_hesvec!,
-               autoback_hesvec, autoback_hesvec!,
-               auto_vecjac, auto_vecjac!,
-               ZygoteVecJac, ZygoteHesVec
+#if !isdefined(Base, :get_extension)
+    using Requires
+#end
 
-        include("differentiation/vecjac_products_zygote.jl")
-        include("differentiation/jaches_products_zygote.jl")
-    end
+function __init__()
+    #@static if !isdefined(Base, :get_extension)
+        @require Zygote="e88e6eb3-aa80-5325-afca-941959d7151f" begin
+            include("../ext/SparseDiffToolsZygote.jl")
+        end
+    #end
 end
 
 end # module