Fix aliasing of x and v in tests, as well as unreverted edits to x in in-place

hesvecs, and update tests accordingly

Author: gaurav-arya

ext/SparseDiffToolsZygote.jl

@@ -25,6 +25,7 @@ function SparseDiffTools.numback_hesvec!(dy, f, x, v, cache1 = similar(v), cache
     g(cache1, x)
     @. x -= 2ϵ * v
     g(cache2, x)
+    @. x += ϵ * v
     @. dy = (cache1 - cache2) / (2ϵ)
 end
 
@@ -37,6 +38,7 @@ function SparseDiffTools.numback_hesvec(f, x, v)
     gxp = g(x)
     x -= 2ϵ * v
     gxm = g(x)
+    @. x += ϵ * v
     (gxp - gxm) / (2ϵ)
 end
 

src/differentiation/jaches_products.jl

@@ -78,6 +78,7 @@ function num_hesvec!(dy,
     g(cache2, x)
     @. x -= 2ϵ * v
     g(cache3, x)
+    @. x += ϵ * v
     @. dy = (cache2 - cache3) / (2ϵ)
 end
 
@@ -110,6 +111,7 @@ function numauto_hesvec!(dy,
     g(cache1, x)
     @. x -= 2ϵ * v
     g(cache2, x)
+    @. x += ϵ * v
     @. dy = (cache1 - cache2) / (2ϵ)
 end
 
@@ -158,6 +160,7 @@ function num_hesvecgrad!(dy, g, x, v, cache2 = similar(v), cache3 = similar(v))
     g(cache2, x)
     @. x -= 2ϵ * v
     g(cache3, x)
+    @. x += ϵ * v
     @. dy = (cache2 - cache3) / (2ϵ)
 end
 

test/test_jaches_products.jl

@@ -4,15 +4,23 @@ using LinearAlgebra, Test
 using Random
 Random.seed!(123)
 N = 300
-const A = rand(N, N)
-
-_f(y, x) = mul!(y, A, x.^2)
-_f(x) = A * (x.^2)
 
 x = rand(N)
 v = rand(N)
+
+# Save original values of x and v to make sure they are not ever mutated
+x0 = copy(x)
+v0 = copy(v)
+
 a, b = rand(2)
 dy = similar(x)
+
+# Define functions for testing
+
+A = rand(N, N)
+_f(y, x) = mul!(y, A, x.^2)
+_f(x) = A * (x.^2)
+
 _g(x) = sum(abs2, x.^2)
 function _h(x)
     FiniteDiff.finite_difference_gradient(_g, x)
@@ -21,7 +29,7 @@ function _h(dy, x)
     FiniteDiff.finite_difference_gradient!(dy, _g, x)
 end
 
-# Define state-dependent functions for operator tests 
+# Make functions state-dependent for operator tests 
 
 include("update_coeffs_testutils.jl")
 f = WrapFunc(_f, 1.0, 1.0)
@@ -47,152 +55,147 @@ cache2 = ForwardDiff.Dual{typeof(ForwardDiff.Tag(SparseDiffTools.DeivVecTag(), e
                   similar(v))≈ForwardDiff.hessian(g, x) * v rtol=1e-2
 @test num_hesvec(g, x, v)≈ForwardDiff.hessian(g, x) * v rtol=1e-2
 
-@test numauto_hesvec!(dy, g, x, v)≈ForwardDiff.hessian(g, x) * v rtol=1e-8
+@test numauto_hesvec!(dy, g, x, v)≈ForwardDiff.hessian(g, x) * v 
 @test numauto_hesvec!(dy, g, x, v, ForwardDiff.GradientConfig(g, x), similar(v),
-                      similar(v))≈ForwardDiff.hessian(g, x) * v rtol=1e-8
-@test numauto_hesvec(g, x, v)≈ForwardDiff.hessian(g, x) * v rtol=1e-8
+                      similar(v))≈ForwardDiff.hessian(g, x) * v 
+@test numauto_hesvec(g, x, v)≈ForwardDiff.hessian(g, x) * v
 
-@test autonum_hesvec!(dy, g, x, v)≈ForwardDiff.hessian(g, x) * v rtol=1e-2
+@test autonum_hesvec!(dy, g, x, v)≈ForwardDiff.hessian(g, x) * v
 @test autonum_hesvec!(dy, g, x, v, cache1, cache2)≈ForwardDiff.hessian(g, x) * v rtol=1e-2
-@test autonum_hesvec(g, x, v)≈ForwardDiff.hessian(g, x) * v rtol=1e-8
+@test autonum_hesvec(g, x, v)≈ForwardDiff.hessian(g, x) * v
 
-@test numback_hesvec!(dy, g, x, v)≈ForwardDiff.hessian(g, x) * v rtol=1e-8
-@test numback_hesvec!(dy, g, x, v, similar(v), similar(v))≈ForwardDiff.hessian(g, x) * v rtol=1e-8
-@test numback_hesvec(g, x, v)≈ForwardDiff.hessian(g, x) * v rtol=1e-8
+@test numback_hesvec!(dy, g, x, v)≈ForwardDiff.hessian(g, x) * v rtol=1e-6
+@test numback_hesvec!(dy, g, x, v, similar(v), similar(v))≈ForwardDiff.hessian(g, x) * v rtol=1e-6
+@test numback_hesvec(g, x, v)≈ForwardDiff.hessian(g, x) * v rtol=1e-6
 
 cache3 = ForwardDiff.Dual{typeof(ForwardDiff.Tag(Nothing, eltype(x))), eltype(x), 1
                           }.(x, ForwardDiff.Partials.(tuple.(v)))
 cache4 = ForwardDiff.Dual{typeof(ForwardDiff.Tag(Nothing, eltype(x))), eltype(x), 1
                           }.(x, ForwardDiff.Partials.(tuple.(v)))
-@test autoback_hesvec!(dy, g, x, v)≈ForwardDiff.hessian(g, x) * v rtol=1e-8
-@test autoback_hesvec!(dy, g, x, v, cache3, cache4)≈ForwardDiff.hessian(g, x) * v rtol=1e-8
-@test autoback_hesvec(g, x, v)≈ForwardDiff.hessian(g, x) * v rtol=1e-8
+@test autoback_hesvec!(dy, g, x, v)≈ForwardDiff.hessian(g, x) * v 
+@test autoback_hesvec!(dy, g, x, v, cache3, cache4)≈ForwardDiff.hessian(g, x) * v
+@test autoback_hesvec(g, x, v)≈ForwardDiff.hessian(g, x) * v
 
 @test num_hesvecgrad!(dy, h, x, v)≈ForwardDiff.hessian(g, x) * v rtol=1e-2
 @test num_hesvecgrad!(dy, h, x, v, similar(v), similar(v))≈ForwardDiff.hessian(g, x) * v rtol=1e-2
 @test num_hesvecgrad(h, x, v)≈ForwardDiff.hessian(g, x) * v rtol=1e-2
 
-@test auto_hesvecgrad!(dy, h, x, v)≈ForwardDiff.hessian(g, x) * v rtol=1e-2
-@test auto_hesvecgrad!(dy, h, x, v, cache1, cache2)≈ForwardDiff.hessian(g, x) * v rtol=1e-2
-@test auto_hesvecgrad(h, x, v)≈ForwardDiff.hessian(g, x) * v rtol=1e-2
+@test auto_hesvecgrad!(dy, h, x, v)≈ForwardDiff.hessian(g, x) * v
+@test auto_hesvecgrad!(dy, h, x, v, cache1, cache2)≈ForwardDiff.hessian(g, x) * v
+@test auto_hesvecgrad(h, x, v)≈ForwardDiff.hessian(g, x) * v
 
 @info "JacVec"
 
-L = JacVec(f, x, 1.0, 1.0)
+L = JacVec(f, copy(x), 1.0, 1.0)
 update_coefficients!(f, x, 1.0, 1.0)
 @test L * x ≈ auto_jacvec(f, x, x)
 @test L * v ≈ auto_jacvec(f, x, v)
 @test mul!(dy, L, v) ≈ auto_jacvec(f, x, v)
 dy=rand(N);_dy=copy(dy);@test mul!(dy,L,v,a,b) ≈ a*auto_jacvec(f,x,v) + b*_dy
 update_coefficients!(L, v, 3.0, 4.0)
 update_coefficients!(f, v, 3.0, 4.0)
-@test mul!(dy, L, v) ≈ auto_jacvec(f, v, v)
-dy=rand(N);_dy=copy(dy);@test mul!(dy,L,v,a,b) ≈ a*auto_jacvec(f,x,v) + b*_dy
+@test mul!(dy, L, x) ≈ auto_jacvec(f, v, x)
+dy=rand(N);_dy=copy(dy);@test mul!(dy,L,x,a,b) ≈ a*auto_jacvec(f,v,x) + b*_dy
 update_coefficients!(f, v, 5.0, 6.0)
-@test L(dy, v, 5.0, 6.0) ≈ auto_jacvec(f,x,v)
+@test L(dy, v, 5.0, 6.0) ≈ auto_jacvec(f,v,v)
 
-L = JacVec(f, x, 1.0, 1.0; autodiff = AutoFiniteDiff())
+L = JacVec(f, copy(x), 1.0, 1.0; autodiff = AutoFiniteDiff())
 update_coefficients!(f, x, 1.0, 1.0)
 @test L * x ≈ num_jacvec(f, x, x)
 @test L * v ≈ num_jacvec(f, x, v)
 @test mul!(dy, L, v)≈num_jacvec(f, x, v) rtol=1e-6
 dy=rand(N);_dy=copy(dy);@test mul!(dy,L,v,a,b) ≈ a*num_jacvec(f,x,v) + b*_dy rtol=1e-6
 update_coefficients!(L, v, 3.0, 4.0)
 update_coefficients!(f, v, 3.0, 4.0)
-@test mul!(dy, L, v)≈num_jacvec(f, v, v) rtol=1e-6
-dy=rand(N);_dy=copy(dy);@test mul!(dy,L,v,a,b) ≈ a*num_jacvec(f,x,v) + b*_dy rtol=1e-6
+@test mul!(dy, L, x)≈num_jacvec(f, v, x) rtol=1e-6
+dy=rand(N);_dy=copy(dy);@test mul!(dy,L,x,a,b) ≈ a*num_jacvec(f,v,x) + b*_dy rtol=1e-6
 update_coefficients!(f, v, 5.0, 6.0)
-@test L(dy, v, 5.0, 6.0) ≈ num_jacvec(f,x,v) rtol=1e-6
+@test L(dy, v, 5.0, 6.0) ≈ num_jacvec(f,v,v) rtol=1e-6
 
 out = similar(v)
 @test_nowarn gmres!(out, L, v)
 
 @info "HesVec"
 
-x = rand(N)
-v = rand(N)
-L = HesVec(g, x, 1.0, 1.0, autodiff = AutoFiniteDiff())
+L = HesVec(g, copy(x), 1.0, 1.0, autodiff = AutoFiniteDiff())
 update_coefficients!(g, x, 1.0, 1.0)
 @test L * x ≈ num_hesvec(g, x, x) rtol=1e-2
-num_hesvec(g, x, x)
 @test L * v ≈ num_hesvec(g, x, v) rtol=1e-2
 @test mul!(dy, L, v)≈num_hesvec(g, x, v) rtol=1e-2
 dy=rand(N);_dy=copy(dy);@test mul!(dy,L,v,a,b) ≈ a*num_hesvec(g,x,v) + b*_dy rtol=1e-2
 update_coefficients!(L, v, 3.0, 4.0)
 update_coefficients!(g, v, 3.0, 4.0)
-@test mul!(dy, L, v)≈num_hesvec(g, v, v) rtol=1e-2
-dy=rand(N);_dy=copy(dy);@test mul!(dy,L,v,a,b) ≈ a*num_hesvec(g,x,v) + b*_dy rtol=1e-2
+@test mul!(dy, L, x)≈num_hesvec(g, v, x) rtol=1e-2
+dy=rand(N);_dy=copy(dy);@test mul!(dy,L,x,a,b) ≈ a*num_hesvec(g,v,x) + b*_dy rtol=1e-2
 update_coefficients!(g, v, 5.0, 6.0)
-@test L(dy, v, 5.0, 6.0) ≈ num_hesvec(g,x,v) rtol=1e-2
+@test L(dy, v, 5.0, 6.0) ≈ num_hesvec(g,v,v) rtol=1e-2
 
-L = HesVec(g, x, 1.0, 1.0)
-update_coefficients!(g, x, 1.0, 1.0)
-numauto_hesvec(g, x, x)
-num_hesvec(g, x, x)
+L = HesVec(g, copy(x), 1.0, 1.0)
 @test L * x ≈ numauto_hesvec(g, x, x)
 @test L * v ≈ numauto_hesvec(g, x, v)
-@test mul!(dy, L, v)≈numauto_hesvec(g, x, v) rtol=1e-8
-dy=rand(N);_dy=copy(dy);@test mul!(dy,L,v,a,b)≈a*numauto_hesvec(g,x,v)+b*_dy rtol=1e-8
+@test mul!(dy, L, v)≈numauto_hesvec(g, x, v)
+dy=rand(N);_dy=copy(dy);@test mul!(dy,L,v,a,0)≈a*numauto_hesvec(g,x,v)+0*_dy
 update_coefficients!(L, v, 3.0, 4.0)
 update_coefficients!(g, v, 3.0, 4.0)
-@test mul!(dy, L, v)≈numauto_hesvec(g, v, v) rtol=1e-8
-dy=rand(N);_dy=copy(dy);@test mul!(dy,L,v,a,b)≈a*numauto_hesvec(g,x,v)+b*_dy rtol=1e-8
+@test mul!(dy, L, x)≈numauto_hesvec(g, v, x) 
+dy=rand(N);_dy=copy(dy);@test mul!(dy,L,x,a,b)≈a*numauto_hesvec(g,v,x)+b*_dy
 update_coefficients!(g, v, 5.0, 6.0)
-@test L(dy, v, 5.0, 6.0) ≈ numauto_hesvec(g,x,v) rtol=1e-2
+@test L(dy, v, 5.0, 6.0) ≈ numauto_hesvec(g,v,v)
 
 out = similar(v)
 gmres!(out, L, v)
 
-x = rand(N)
-v = rand(N)
-
-L = HesVec(g, x, 1.0, 1.0; autodiff = AutoZygote())
+L = HesVec(g, copy(x), 1.0, 1.0; autodiff = AutoZygote())
 update_coefficients!(g, x, 1.0, 1.0)
 @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
+@test mul!(dy, L, v)≈autoback_hesvec(g, x, v)
+dy=rand(N);_dy=copy(dy);@test mul!(dy,L,v,a,b)≈a*autoback_hesvec(g,x,v)+b*_dy
 update_coefficients!(L, v, 3.0, 4.0)
 update_coefficients!(g, v, 3.0, 4.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
+@test mul!(dy, L, x)≈autoback_hesvec(g, v, x)
+dy=rand(N);_dy=copy(dy);@test mul!(dy,L,x,a,b)≈a*autoback_hesvec(g,v,x)+b*_dy
 update_coefficients!(g, v, 5.0, 6.0)
-@test L(dy, v, 5.0, 6.0) ≈ autoback_hesvec(g,x,v) rtol=1e-2
+@test L(dy, v, 5.0, 6.0) ≈ autoback_hesvec(g,v,v)
 
 out = similar(v)
 gmres!(out, L, v)
 
 @info "HesVecGrad"
 
-x = rand(N)
-v = rand(N)
-L = HesVecGrad(h, x, 1.0, 1.0; autodiff = AutoFiniteDiff())
+L = HesVecGrad(h, copy(x), 1.0, 1.0; autodiff = AutoFiniteDiff())
 update_coefficients!(h, x, 1.0, 1.0)
 update_coefficients!(g, x, 1.0, 1.0)
 @test L * x ≈ num_hesvec(g, x, x) rtol=1e-2
 @test L * v ≈ num_hesvec(g, x, v) rtol=1e-2
 @test mul!(dy, L, v)≈num_hesvec(g, x, v) rtol=1e-2
 dy=rand(N);_dy=copy(dy);@test mul!(dy,L,v,a,b)≈a*num_hesvec(g,x,v)+b*_dy rtol=1e-2
 for op in (L, g, h) update_coefficients!(op, v, 3.0, 4.0) end
-@test mul!(dy, L, v)≈num_hesvec(g, v, v) rtol=1e-2
-dy=rand(N);_dy=copy(dy);@test mul!(dy,L,v,a,b)≈a*num_hesvec(g,x,v)+b*_dy rtol=1e-2
+@test mul!(dy, L, x)≈num_hesvec(g, v, x) rtol=1e-2
+dy=rand(N);_dy=copy(dy);@test mul!(dy,L,x,a,b)≈a*num_hesvec(g,v,x)+b*_dy rtol=1e-2
 update_coefficients!(g, v, 5.0, 6.0)
-@test L(dy, v, 5.0, 6.0) ≈ num_hesvec(g,x,v) rtol=1e-2
+@test L(dy, v, 5.0, 6.0) ≈ num_hesvec(g,v,v) rtol=1e-2
 
-L = HesVecGrad(h, x, 1.0, 1.0)
+L = HesVecGrad(h, copy(x), 1.0, 1.0)
 update_coefficients!(g, x, 1.0, 1.0)
 update_coefficients!(h, x, 1.0, 1.0)
 @test L * x ≈ autonum_hesvec(g, x, x)
 @test L * v ≈ numauto_hesvec(g, x, v)
 @test mul!(dy, L, v)≈numauto_hesvec(g, x, v) rtol=1e-8
-dy=rand(N);_dy=copy(dy);@test mul!(dy,L,v,a,b)≈a*numauto_hesvec(g,x,v)+b*_dy rtol=1e-8
+dy=rand(N);_dy=copy(dy);@test mul!(dy,L,v,a,b)≈a*numauto_hesvec(g,x,v)+b*_dy
 for op in (L, g, h) update_coefficients!(op, v, 3.0, 4.0) end
-@test mul!(dy, L, v)≈numauto_hesvec(g, v, v) rtol=1e-8
-dy=rand(N);_dy=copy(dy);@test mul!(dy,L,v,a,b)≈a*numauto_hesvec(g,x,v)+b*_dy rtol=1e-8
+@test mul!(dy, L, x)≈numauto_hesvec(g, v, x)
+dy=rand(N);_dy=copy(dy);@test mul!(dy,L,x,a,b)≈a*numauto_hesvec(g,v,x)+b*_dy
 update_coefficients!(g, v, 5.0, 6.0)
 update_coefficients!(h, v, 5.0, 6.0)
-@test L(dy, v, 5.0, 6.0) ≈ num_hesvec(g,x,v) rtol=1e-2
+@test L(dy, v, 5.0, 6.0) ≈ numauto_hesvec(g,v,v)
 
 out = similar(v)
 gmres!(out, L, v)
+
+# Test that x and v were not mutated
+# x's rtol can't be too large since it is mutated and then restored in some algorithms
+@test x ≈ x0
+@test v ≈ v0
+
 #

test/test_vecjac_products.jl

@@ -4,33 +4,61 @@ using LinearAlgebra, Test
 using Random
 Random.seed!(123)
 N = 300
-const A = rand(N, N)
 
+# Use Float32 since Zygote defaults to Float32
 x = 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)
+
+a, b = rand(2)
+dy = similar(x)
+
+A = rand(Float32, N, N)
 _f(du,u) = mul!(du, A, u)
 _f(u) = A * u
 
 # Define state-dependent functions for operator tests 
 include("update_coeffs_testutils.jl")
-f = WrapFunc(_f, 1.0, 1.0)
+f = WrapFunc(_f, 1.0f0, 1.0f0)
+
+# Compute Jacobian via Zygote
 
 @info "VecJac"
 
-L = VecJac(f, x, 1.0, 1.0)
+L = VecJac(f, copy(x), 1.0f0, 1.0f0; autodiff = AutoZygote())
+update_coefficients!(f, x, 1.0, 1.0)
+actual_jac = Zygote.jacobian(f, x)[1]
+@test L * x ≈ actual_jac' * x 
+@test L * v ≈ actual_jac' * v 
+@test mul!(dy, L, v) ≈ actual_jac' * v 
 update_coefficients!(L, v, 3.0, 4.0)
 update_coefficients!(f, v, 3.0, 4.0)
-actual_vjp = Zygote.jacobian(f, x)[1]' * v
-@test L * v ≈ actual_vjp
+actual_jac = Zygote.jacobian(f, v)[1]
+@test mul!(dy, L, x) ≈ actual_jac' * x
+_dy=copy(dy); @test mul!(dy,L,x,a,b) ≈ a*actual_jac'*x + b*_dy
 update_coefficients!(f, v, 5.0, 6.0)
-actual_vjp2 = Zygote.jacobian(f, x)[1]' * v
-@test L(copy(v), v, 5.0, 6.0) ≈ actual_vjp2
+actual_jac = Zygote.jacobian(f, v)[1]
+@test L(dy, v, 5.0, 6.0) ≈ actual_jac' * v
 
-L = VecJac(f, x, 1.0, 1.0; autodiff = AutoFiniteDiff())
+L = VecJac(f, copy(x), 1.0f0, 1.0f0; autodiff = AutoFiniteDiff())
+update_coefficients!(f, x, 1.0, 1.0)
+actual_jac = Zygote.jacobian(f, x)[1]
+@test L * x ≈ actual_jac' * x 
+@test L * v ≈ actual_jac' * v 
+@test mul!(dy, L, v) ≈ actual_jac' * v 
 update_coefficients!(L, v, 3.0, 4.0)
 update_coefficients!(f, v, 3.0, 4.0)
-@test L * v ≈ actual_vjp
+actual_jac = Zygote.jacobian(f, v)[1]
+@test mul!(dy, L, x) ≈ actual_jac' * x
+_dy=copy(dy); @test mul!(dy,L,x,a,b) ≈ a*actual_jac'*x + b*_dy
 update_coefficients!(f, v, 5.0, 6.0)
-@test L(copy(v), v, 5.0, 6.0) ≈ actual_vjp2
+actual_jac = Zygote.jacobian(f, v)[1]
+@test L(dy, v, 5.0, 6.0) ≈ actual_jac' * v
+
+# Test that x and v were not mutated
+@test x ≈ x0
+@test v ≈ v0
 #