zygote

Author: vpuri3

src/SparseDiffTools.jl

@@ -71,8 +71,10 @@ 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!
+        export numback_hesvec, numback_hesvec!,
+               autoback_hesvec, autoback_hesvec!,
+               auto_vecjac, auto_vecjac!,
+               ZygoteVecJac, ZygoteHesVec
 
         include("differentiation/vecjac_products_zygote.jl")
         include("differentiation/jaches_products_zygote.jl")

src/differentiation/jaches_products_zygote.jl

@@ -25,21 +25,21 @@ function numback_hesvec(f, x, v)
 end
 
 function autoback_hesvec!(dy, f, x, v,
-                          cache2 = Dual{typeof(ForwardDiff.Tag(DeivVecTag, eltype(x))),
+                          cache1 = Dual{typeof(ForwardDiff.Tag(DeivVecTag, eltype(x))),
                                         eltype(x), 1
                                         }.(x,
                                            ForwardDiff.Partials.(Tuple.(reshape(v, size(x))))),
-                          cache3 = Dual{typeof(ForwardDiff.Tag(DeivVecTag, eltype(x))),
+                          cache2 = Dual{typeof(ForwardDiff.Tag(DeivVecTag, eltype(x))),
                                         eltype(x), 1
                                         }.(x,
                                            ForwardDiff.Partials.(Tuple.(reshape(v, size(x))))))
     g = let f = f
         (dx, x) -> dx .= first(Zygote.gradient(f, x))
     end
-    cache2 .= Dual{typeof(ForwardDiff.Tag(DeivVecTag, eltype(x))), eltype(x), 1
+    cache1 .= Dual{typeof(ForwardDiff.Tag(DeivVecTag, eltype(x))), eltype(x), 1
                    }.(x, ForwardDiff.Partials.(Tuple.(reshape(v, size(x)))))
-    g(cache3, cache2)
-    dy .= partials.(cache3, 1)
+    g(cache2, cache1)
+    dy .= partials.(cache2, 1)
 end
 
 function autoback_hesvec(f, x, v)
@@ -48,3 +48,38 @@ function autoback_hesvec(f, x, v)
              }.(x, ForwardDiff.Partials.(Tuple.(reshape(v, size(x)))))
     ForwardDiff.partials.(g(y), 1)
 end
+
+### Operator Forms
+
+function ZygoteHesVec(f, u::AbstractArray, p = nothing, t = nothing; autodiff = true)
+
+    if autodiff
+        cache1 = Dual{
+                      typeof(ForwardDiff.Tag(DeivVecTag(),eltype(u))), eltype(u), 1
+                     }.(u, ForwardDiff.Partials.(tuple.(u)))
+        cache2 = copy(u)
+    else
+        cache1 = similar(u)
+        cache2 = similar(u)
+    end
+
+    cache = (cache1, cache2,)
+
+    vecprod  = autodiff ? autoback_hesvec  : numback_hesvec 
+    vecprod! = autodiff ? autoback_hesvec! : numback_hesvec!
+
+    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
+#

src/differentiation/vecjac_products_zygote.jl

@@ -7,3 +7,6 @@ function auto_vecjac(f, x, v)
     vv, back = Zygote.pullback(f, x)
     return vec(back(reshape(v, size(vv)))[1])
 end
+
+const ZygoteVecJac = VecJac
+#

test/test_ad.jl

@@ -4,8 +4,7 @@ using SparseArrays, Test
 using LinearAlgebra
 using BlockBandedMatrices, ArrayInterfaceBlockBandedMatrices
 using BandedMatrices, ArrayInterfaceBandedMatrices
-using StaticArrays
-using ArrayInterfaceStaticArrays
+using StaticArrays#, StaticArrayInterface
 
 fcalls = 0
 function f(dx, x)

test/test_jaches_products.jl

@@ -113,6 +113,33 @@ 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"
+using Zygote
+x = rand(N)
+v = rand(N)
+
+L = ZygoteHesVec(g, x, autodiff = false)
+@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
+dy=rand(N);_dy=copy(dy);@test mul!(dy,L,v,a,b) ≈ a*numback_hesvec(g,x,v) + b*_dy rtol=1e-2
+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 = HesVec(g, x)
+@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
+dy=rand(N);_dy=copy(dy);@test mul!(dy,L,v,a,b)≈a*autoback_hesvec(g,x,v)+b*_dy rtol=1e-8
+update_coefficients!(L, v, nothing, 0.0)
+@test mul!(dy, L, v)≈autoback_hesvec(g, v, v) rtol=1e-8
+dy=rand(N);_dy=copy(dy);@test mul!(dy,L,v,a,b)≈a*autoback_hesvec(g,x,v)+b*_dy rtol=1e-8
+
+out = similar(v)
+gmres!(out, L, v)
+
+
 @info "HesVecGrad"
 
 x = rand(N)

test/test_vecjac_products.jl

@@ -12,7 +12,7 @@ v = rand(Float32, N)
 f(du,u,p,t) = mul!(du, A, u)
 f(u,p,t) = A * u
 
-# VecJac
+@info "VecJac"
 
 L = VecJac(f, x)
 actual_vjp = Zygote.jacobian(x -> f(x, nothing, 0.0), x)[1]' * v
@@ -21,4 +21,14 @@ update_coefficients!(L, v, nothing, 0.0)
 L = VecJac(f, x; autodiff = false)
 update_coefficients!(L, v, nothing, 0.0)
 @test L * v ≈ actual_vjp
+
+@info "ZygoteVecJac"
+
+L = ZygoteVecJac(f, x)
+actual_vjp = Zygote.jacobian(x -> f(x, nothing, 0.0), x)[1]' * v
+update_coefficients!(L, v, nothing, 0.0)
+@test L * v ≈ actual_vjp
+L = ZygoteVecJac(f, x; autodiff = false)
+update_coefficients!(L, v, nothing, 0.0)
+@test L * v ≈ actual_vjp
 #