Zygote

Author: vpuri3

src/differentiation/vecjac_products.jl

@@ -86,9 +86,13 @@ function (L::RevModeAutoDiffVecProd{ad,true,false})(dv, v, p, t) where{ad}
     L.vecprod!(dv, (_du, _u) -> L.f(_du, _u, p, t), L.u, v, L.cache...)
 end
 
-function VecJac(f, u::AbstractArray, p = nothing, t = nothing; autodiff = true,
+function VecJac(f, u::AbstractArray, p = nothing, t = nothing; autodiff = false,
                 ishermitian = false, opnrom = true)
 
+    if autodiff
+        @assert isdefined(SparseDiffTools, :auto_vecjac) "Please load Zygote with `using Zygote`, or `import Zygote` to use VecJac with `autodiff = true`."
+    end
+
     cache = (similar(u), similar(u),)
 
     vecprod  = autodiff ? auto_vecjac  : num_vecjac

src/differentiation/vecjac_products_zygote.jl

@@ -8,5 +8,7 @@ function auto_vecjac(f, x, v)
     return vec(back(reshape(v, size(vv)))[1])
 end
 
-const ZygoteVecJac = VecJac
+#ZygoteVecJac = VecJac
+ZygoteVecJac(args...; autodiff = true, kwargs...) = VecJac(args...; autodiff = autodiff, kwargs...)
+
 #