Zygote ext

Project.toml

@@ -1,7 +1,7 @@
 name = "SparseDiffTools"
 uuid = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
 authors = ["Pankaj Mishra <[email protected]>", "Chris Rackauckas <[email protected]>"]
-version = "2.00.0"
+version = "2.0.0"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -19,6 +19,12 @@ 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]
 Adapt = "3.0"
 ArrayInterface = "6, 7"

src/differentiation/jaches_products_zygote.jl → 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,15 +68,15 @@ 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")
+function __init__()
+    @static if !isdefined(Base, :get_extension)
+        @require Zygote="e88e6eb3-aa80-5325-afca-941959d7151f" begin
+            include("../ext/SparseDiffToolsZygote.jl")
+        end
     end
 end