use kwarg VJP_input for recomputing pullback

vpuri3 · vpuri3 · commit 46d5f6498f5b · 2023-05-30T21:17:56.000-04:00
diff --git a/ext/SparseDiffToolsZygote.jl b/ext/SparseDiffToolsZygote.jl
@@ -91,21 +91,33 @@ function SparseDiffTools._vecjac(f, u, autodiff::AutoZygote)
     AutoDiffVJP(f, u, cache, autodiff, pullback)
 end
 
-function update_coefficients(L::AutoDiffVJP{AD}, u, p, t) where{AD <: AutoZygote}
-    @set! L.f = update_coefficients(L.f, u, p, t)
-    @set! L.u = u
-    @set! L.pullback = Zygote.pullback(L.f, u)
+function update_coefficients(L::AutoDiffVJP{AD}, u, p, t; VJP_input = nothing,
+                            ) where{AD <: AutoZygote}
+
+    if !isnothing(VJP_input)
+        @set! L.u = VJP_input
+    end
+
+    @set! L.f = update_coefficients(L.f, L.u, p, t)
+    @set! L.pullback = Zygote.pullback(L.f, L.u)
 end
 
-function update_coefficients!(L::AutoDiffVJP{AD}, u, p, t) where{AD <: AutoZygote}
-    update_coefficients!(L.f, u, p, t)
-    copy!(L.u, u)
-    L.pullback = Zygote.pullback(L.f, u)
+function update_coefficients!(L::AutoDiffVJP{AD}, u, p, t; VJP_input = nothing,
+                             ) where{AD <: AutoZygote}
+
+    if !isnothing(VJP_input)
+        copy!(L.u, VJP_input)
+    end
+
+    update_coefficients!(L.f, L.u, p, t)
+    L.pullback = Zygote.pullback(L.f, L.u)
+
     L
 end
 
 # Interpret the call as df/du' * v
-function (L::AutoDiffVJP{AD})(v, p, t) where{AD <: AutoZygote}
+function (L::AutoDiffVJP{AD})(v, p, t; VJP_input = nothing) where{AD <: AutoZygote}
+    # ignore VJP_input as pullback was computed in update_coefficients(...)
 
     y, back = L.pullback
     V = reshape(v, size(y))
@@ -114,13 +126,15 @@ function (L::AutoDiffVJP{AD})(v, p, t) where{AD <: AutoZygote}
 end
 
 # prefer non in-place method
-function (L::AutoDiffVJP{AD, IIP, true})(dv, v, p, t) where {AD <: AutoZygote, IIP}
+function (L::AutoDiffVJP{AD, IIP, true})(dv, v, p, t; VJP_input = nothing) where {AD <: AutoZygote, IIP}
+    # ignore VJP_input as pullback was computed in update_coefficients!(...)
+
     _dv = L(v, p, t)
     copy!(dv, _dv)
 end
 
-function (L::AutoDiffVJP{AD, true, false})(dv, v, p, t) where {AD <: AutoZygote}
-    SparseDiffTools.auto_vecjac!(dv, L.f, L.u, v, L.cache...)
+function (L::AutoDiffVJP{AD, true, false})(dv, v, p, t; VJP_input = nothing) where {AD <: AutoZygote}
+    @error("Zygote requires an out of place method with signature f(u).")
 end
 
 end # module
diff --git a/src/differentiation/vecjac_products.jl b/src/differentiation/vecjac_products.jl
@@ -39,15 +39,34 @@ end
 
 """
     VecJac(f, u, [p, t]; autodiff = AutoFiniteDiff())
+
+Returns SciMLOperators.FunctionOperator which computes vector-jacobian
+product `df/du * v`.
+
+```
+L = VecJac(f, u)
+
+L * v         # = df/du * v
+mul!(w, L, v) # = df/du * v
+
+L(v, p, t; VJP_input = w)    # = df/dw * v
+L(x, v, p, t; VJP_input = w) # = df/dw * v
+```
 """
 function VecJac(f, u::AbstractArray, p = nothing, t = nothing;
                 autodiff = AutoFiniteDiff(), kwargs...)
 
     L = _vecjac(f, u, autodiff)
     IIP, OOP = get_iip_oop(L)
 
+    if isa(autodiff, AutoZygote) & !OOP
+        msg = "Zygote requires an out of place method with signature f(u)."
+        throw(ArgumentError(msg))
+    end
+
     FunctionOperator(L, u, u; isinplace = IIP, outofplace = OOP,
-                     p = p, t = t, islinear = true, kwargs...)
+                     p = p, t = t, islinear = true,
+                     accepted_kwargs = (:VJP_input,), kwargs...)
 end
 
 function _vecjac(f, u, autodiff::AutoFiniteDiff)
@@ -59,10 +78,15 @@ function _vecjac(f, u, autodiff::AutoFiniteDiff)
 end
 
 mutable struct AutoDiffVJP{AD, IIP, OOP, F, U, C, PB} <: AbstractAutoDiffVecProd
+    """ Compute VJP of `f` at `u`, applied to vector `v`: `df/du' * u` """
     f::F
+    """ input to `f` """
     u::U
+    """ Cache for num_vecjac! when autodiff isa AutoFintieDiff """
     cache::C
+    """ Type of automatic differentiation algorithm """
     autodiff::AD
+    """ stores the result of Zygote.pullback for AutoZygote """
     pullback::PB
 
     function AutoDiffVJP(f, u, cache, autodiff, pullback)
@@ -93,23 +117,36 @@ function get_iip_oop(::AutoDiffVJP{AD, IIP, OOP}) where{AD, IIP, OOP}
     IIP, OOP
 end
 
-function update_coefficients(L::AutoDiffVJP{AD}, u, p, t) where{AD <: AutoFiniteDiff}
-    @set! L.f = update_coefficients(L.f, u, p, t)
-    @set! L.u = u
+function update_coefficients(L::AutoDiffVJP{AD}, u, p, t; VJP_input = nothing,
+                            ) where{AD <: AutoFiniteDiff}
+
+    if !isnothing(VJP_input)
+        @set! L.u = VJP_input
+    end
+
+    @set! L.f = update_coefficients(L.f, L.u, p, t)
 end
 
-function update_coefficients!(L::AutoDiffVJP{AD}, u, p, t) where{AD <: AutoFiniteDiff}
-    update_coefficients!(L.f, u, p, t)
-    copy!(L.u, u)
+function update_coefficients!(L::AutoDiffVJP{AD}, u, p, t; VJP_input = nothing,
+                             ) where{AD <: AutoFiniteDiff}
+
+    if !isnothing(VJP_input)
+        copy!(L.u, VJP_input)
+    end
+
+    update_coefficients!(L.f, L.u, p, t)
+
     L
 end
 
 # Interpret the call as df/du' * v
-function (L::AutoDiffVJP{AD})(v, p, t) where{AD <: AutoFiniteDiff}
+function (L::AutoDiffVJP{AD})(v, p, t; VJP_input = nothing,) where{AD <: AutoFiniteDiff}
+    # ignore VJP_input as L.u was set in update_coefficients(...)
     num_vecjac(L.f, L.u, v)
 end
 
-function (L::AutoDiffVJP{AD})(dv, v, p, t) where{AD <: AutoFiniteDiff}
+function (L::AutoDiffVJP{AD})(dv, v, p, t; VJP_input = nothing,) where{AD <: AutoFiniteDiff}
+    # ignore VJP_input as L.u was set in update_coefficients!(...)
     num_vecjac!(dv, L.f, L.u, v, L.cache...)
 end
 
diff --git a/test/test_vecjac_products.jl b/test/test_vecjac_products.jl
@@ -6,12 +6,15 @@ Random.seed!(123)
 N = 300
 
 # Use Float32 since Zygote defaults to Float32
-x = rand(Float32, N)
+x1 = rand(Float32, N)
+x2 = rand(Float32, N)
+
 v = rand(Float32, N)
 
 # Save original values of x and v to make sure they are not ever mutated
-x0 = copy(x)
-v0 = copy(v)
+_x1 = copy(x1)
+_x2 = copy(x2)
+_v = copy(v)
 
 a, b = rand(Float32, 2)
 
@@ -23,81 +26,106 @@ _f(x) = A * (x .^ 2)
 include("update_coeffs_testutils.jl")
 f = WrapFunc(_f, 1.0f0, 1.0f0)
 
-@test auto_vecjac(f, x, v) ≈ Zygote.jacobian(f, x)[1]' * v
-@test auto_vecjac!(zero(x), f, x, v) ≈ auto_vecjac(f, x, v)
-@test num_vecjac!(zero(x), f, copy(x), v) ≈ num_vecjac(f, copy(x), v)
-@test auto_vecjac(f, x, v) ≈ num_vecjac(f, copy(x), copy(v)) rtol = 1e-2
+@test auto_vecjac(f, x1, v) ≈ Zygote.jacobian(f, x1)[1]' * v
+@test auto_vecjac!(zero(x1), f, x1, v) ≈ auto_vecjac(f, x1, v)
+@test num_vecjac!(zero(x1), f, copy(x1), v) ≈ num_vecjac(f, copy(x1), v)
+@test auto_vecjac(f, x1, v) ≈ num_vecjac(f, copy(x1), copy(v)) rtol = 1e-2
 
 # Compute Jacobian via Zygote
 
 @info "VecJac AutoZygote"
 
-L = VecJac(f, copy(x), 1.0f0, 1.0f0; autodiff = AutoZygote())
+p, t = rand(Float32, 2)
+L = VecJac(f, copy(x1), p, t; autodiff = AutoZygote())
+update_coefficients!(L, v, p, t)
 
-Jx = Zygote.jacobian(f, x)[1]
-Jv = Zygote.jacobian(f, v)[1]
+update_coefficients!(f, v, p, t)
+J1 = Zygote.jacobian(f, x1)[1]
+J2 = Zygote.jacobian(f, x2)[1]
 
-@test L * x ≈ Jx' * x
-@test L * v ≈ Jx' * v
-y=zero(x); @test mul!(y, L, v) ≈ Jx' * v
-y=zero(x); @test mul!(y, L, v) ≈ Jx' * v
+# test operator application
+@test L * v ≈ J1' * v
+@test L(v, p, t) ≈ J1' * v
+y=zeros(N); @test mul!(y, L, v) ≈ J1' * v
+y=zeros(N); @test L(y, v, p, t) ≈ J1' * v
 
-@test L(x, 1.0f0, 1.0f0) ≈ Jx' * x
-y=zero(x); @test L(y, x, 1.0f0, 1.0f0) ≈ Jx' * x
-@test L(v, 1.0f0, 1.0f0) ≈ Jv' * v
-y=zero(v); @test L(y, v, 1.0f0, 1.0f0) ≈ Jv' * v
+# use kwarg VJP_input = x2
+@test L(v, p, t; VJP_input = x2) ≈ J2' * v
+y=zeros(N); @test L(y, v, p, t; VJP_input = x2) ≈ J2' * v
 
-update_coefficients!(L, v, 3.0, 4.0)
+# update_coefficients
+p, t = rand(Float32, 2)
+L = update_coefficients(L, v, p, t; JVP_input = x2)
 
-Jx = Zygote.jacobian(f, x)[1]
-Jv = Zygote.jacobian(f, v)[1]
+update_coefficients!(f, v, p, t)
+J1 = Zygote.jacobian(f, x1)[1]
+J2 = Zygote.jacobian(f, x2)[1]
 
-@test L * x ≈ Jv' * x
-@test L * v ≈ Jv' * v
-y=zero(x); @test mul!(y, L, v) ≈ Jv' * v
-y=zero(x); @test mul!(y, L, v) ≈ Jv' * v
+# @show p, t
+# @show f.p, f.t
+# @show L.op.f.p, L.op.f.t
 
-@test L(x, 3.0f0, 4.0f0) ≈ Jx' * x
-y=zero(x); @test L(y, x, 3.0f0, 4.0f0) ≈ Jx' * x
-@test L(v, 3.0f0, 4.0f0) ≈ Jv' * v
-y=zero(v); @test L(y, v, 3.0f0, 4.0f0) ≈ Jv' * v
+@test L * v ≈ J2' * v
+@test L(v, p, t) ≈ J2' * v
+y=zeros(N); @test mul!(y, L, v) ≈ J2' * v
+y=zeros(N); @test L(y, v, p, t) ≈ J2' * v
+
+# use kwarg VJP_input = x1
+@test L(v, p, t; VJP_input = x1) ≈ J1' * v
+y=zeros(N); @test L(y, v, p, t; VJP_input = x1) ≈ J1' * v
 
 @info "VecJac AutoFiniteDiff"
 
-L = VecJac(f, copy(x), 1.0f0, 1.0f0; autodiff = AutoFiniteDiff())
+p, t = rand(Float32, 2)
+L = VecJac(f, copy(x1), 1.0f0, 1.0f0; autodiff = AutoFiniteDiff())
+update_coefficients!(L, v, p, t)
+update_coefficients!(f, v, p, t)
+
+@test L * v ≈ num_vecjac(f, copy(x1), v)
+@test L(v, p, t) ≈ num_vecjac(f, copy(x1), v)
+y=zeros(N); @test mul!(y, L, v) ≈ num_vecjac(f, copy(x1), v)
+y=zeros(N); @test L(y, v, p, t) ≈ num_vecjac(f, copy(x1), v)
 
-@test L * x ≈ num_vecjac(f, copy(x), x)
-@test L * v ≈ num_vecjac(f, copy(x), v)
-y=zero(x); @test mul!(y, L, v) ≈ num_vecjac(f, copy(x), v)
+# use kwarg VJP_input = x2
+@test L(v, p, t; VJP_input = x2) ≈ num_vecjac(f, copy(x2), v)
+y=zeros(N); @test L(y, v, p, t; VJP_input = x2) ≈ num_vecjac(f, copy(x2), v)
 
-update_coefficients!(L, v, 3.0, 4.0)
-@test mul!(y, L, x) ≈ num_vecjac(f, copy(v), x)
-_y = copy(y); @test mul!(y, L, x, a, b) ≈ a * num_vecjac(f,copy(v),x) + b * _y
+# update_coefficients
+p, t = rand(Float32, 2)
+L = update_coefficients(L, v, p, t; JVP_input = x2)
+update_coefficients!(f, v, p, t)
 
-update_coefficients!(f, v, 5.0, 6.0)
-@test L(y, v, 5.0, 6.0) ≈ num_vecjac(f, copy(v), v)
+@test L * v ≈ num_vecjac(f, copy(x2), v)
+@test L(v, p, t) ≈ num_vecjac(f, copy(x2), v)
+y=zeros(N); @test mul!(y, L, v) ≈ num_vecjac(f, copy(x2), v)
+y=zeros(N); @test L(y, v, p, t) ≈ num_vecjac(f, copy(x2), v)
+
+# use kwarg VJP_input = x2
+@test L(v, p, t; VJP_input = x1) ≈ num_vecjac(f, copy(x1), v)
+y=zeros(N); @test L(y, v, p, t; VJP_input = x1) ≈ num_vecjac(f, copy(x1), v)
 
 # Test that x and v were not mutated
-@test x ≈ x0
-@test v ≈ v0
+@test x1 ≈ _x1
+@test x2 ≈ _x2
+@test v ≈ v
 
 @info "Base.resize!"
 
 # Resize test
 f2(x) = 2x
 f2(y, x) = (copy!(y, x); lmul!(2, y); y)
 
+x = rand(Float32, N)
 for M in (100, 400)
     local L = VecJac(f2, copy(x), 1.0f0, 1.0f0; autodiff = AutoZygote())
     resize!(L, M)
 
     _x = resize!(copy(x), M)
     _u = rand(M)
-    J2 = Zygote.jacobian(f2, _x)[1]
+    local J2 = Zygote.jacobian(f2, _x)[1]
 
-    update_coefficients!(L, _x, 1.0f0, 1.0f0)
+    update_coefficients!(L, _u, 1.0f0, 1.0f0; VJP_input = _x)
     @test L * _u ≈ J2' * _u rtol=1e-6
-    _v = zeros(M); @test mul!(_v, L, _u) ≈ J2' * _u rtol=1e-6
+    local _v = zeros(M); @test mul!(_v, L, _u) ≈ J2' * _u rtol=1e-6
 end
-
 #
