Change rrule verbosity behaviour

Project.toml

@@ -27,7 +27,7 @@ Printf = "1"
 Random = "1"
 Test = "1"
 TestExtras = "0.2,0.3"
-VectorInterface = "0.4"
+VectorInterface = "0.4,0.5"
 Zygote = "0.6"
 julia = "1.6"
 

ext/KrylovKitChainRulesCoreExt/eigsolve.jl

@@ -135,7 +135,7 @@ function compute_eigsolve_pullback_data(Δvals, Δvecs, vals, vecs, info, which,
             b = (zerovector(v), convert(T, Δλ))
         else
             vdΔv = inner(v, Δv)
-            if alg_rrule.verbosity >= 0
+            if alg_primal.verbosity >= 1
                 gauge = abs(imag(vdΔv))
                 gauge > alg_primal.tol &&
                     @warn "`eigsolve` cotangent for eigenvector $i is sensitive to gauge choice: (|gauge| = $gauge)"
@@ -152,9 +152,9 @@ function compute_eigsolve_pullback_data(Δvals, Δvecs, vals, vecs, info, which,
                 return (y1, y2)
             end
         end
-        if info.converged >= i && reverse_info.converged == 0 && alg_rrule.verbosity >= 0
+        if info.converged >= i && reverse_info.converged == 0 && alg_primal.verbosity >= 1
             @warn "`eigsolve` cotangent linear problem ($i) did not converge, whereas the primal eigenvalue problem did: normres = $(reverse_info.normres)"
-        elseif abs(w[2]) > alg_rrule.tol && alg_rrule.verbosity >= 0
+        elseif abs(w[2]) > alg_rrule.tol && alg_primal.verbosity >= 1
             @warn "`eigsolve` cotangent linear problem ($i) returns unexpected result: error = $(w[2])"
         end
         ws[i] = w[1]
@@ -185,7 +185,7 @@ function compute_eigsolve_pullback_data(Δvals, Δvecs, vals, vecs, info, which,
 
     # components along subspace spanned by current eigenvectors
     tol = alg_primal.tol
-    if alg_rrule.verbosity >= 0
+    if alg_primal.verbosity >= 1
         mask = abs.(transpose(vals) .- vals) .< tol
         gaugepart = VdΔV[mask] - Diagonal(real(diag(VdΔV)))[mask]
         Δgauge = norm(gaugepart, Inf)
@@ -263,7 +263,7 @@ function compute_eigsolve_pullback_data(Δvals, Δvecs, vals, vecs, info, which,
             return (w′, conj.(vals) .* x)
         end
     end
-    if info.converged >= n && reverse_info.converged < n && alg_rrule.verbosity >= 0
+    if info.converged >= n && reverse_info.converged < n && alg_primal.verbosity >= 1
         @warn "`eigsolve` cotangent problem did not converge, whereas the primal eigenvalue problem did"
     end
     # cleanup and construct final result by renormalising the eigenvectors and explicitly
@@ -276,7 +276,7 @@ function compute_eigsolve_pullback_data(Δvals, Δvecs, vals, vecs, info, which,
         w, x = Ws[ic]
         factor = 1 / x[i]
         x[i] = zero(x[i])
-        if alg_rrule.verbosity >= 0
+        if alg_primal.verbosity >= 1
             error = max(norm(x, Inf), abs(rvals[ic] - conj(vals[i])))
             error > 10 * tol &&
                 @warn "`eigsolve` cotangent linear problem ($i) returns unexpected result: error = $error"
@@ -308,7 +308,7 @@ function compute_eigsolve_pullback_data(Δvals, Δvecs, vals, vecs, info, which,
     # components along subspace spanned by current eigenvectors
     tol = alg_primal.tol
     aVdΔV = rmul!(VdΔV - VdΔV', 1 / 2)
-    if alg_rrule.verbosity >= 0
+    if alg_primal.verbosity >= 1
         mask = abs.(transpose(vals) .- vals) .< tol
         gaugepart = view(aVdΔV, mask)
         gauge = norm(gaugepart, Inf)
@@ -366,7 +366,7 @@ function compute_eigsolve_pullback_data(Δvals, Δvecs, vals, vecs, info, which,
             return (w′, vals .* x)
         end
     end
-    if info.converged >= n && reverse_info.converged < n && alg_rrule.verbosity >= 0
+    if info.converged >= n && reverse_info.converged < n && alg_primal.verbosity >= 1
         @warn "`eigsolve` cotangent problem did not converge, whereas the primal eigenvalue problem did"
     end
 
@@ -380,7 +380,7 @@ function compute_eigsolve_pullback_data(Δvals, Δvecs, vals, vecs, info, which,
         factor = 1 / x[ic]
         x[ic] = zero(x[ic])
         error = max(norm(x, Inf), abs(rvals[i] - conj(vals[ic])))
-        if error > 5 * tol && alg_rrule.verbosity >= 0
+        if error > 10 * tol && alg_primal.verbosity >= 1
             @warn "`eigsolve` cotangent linear problem ($ic) returns unexpected result: error = $error"
         end
         ws[ic] = VectorInterface.add!!(zs[ic], Q(w), -factor)

ext/KrylovKitChainRulesCoreExt/linsolve.jl

@@ -34,7 +34,7 @@ function make_linsolve_pullback(fᴴ, b, a₀, a₁, alg_rrule, construct∂f, x
                                                                                      a₁)))
         ∂b, reverse_info = linsolve(fᴴ, x̄, x̄₀, alg_rrule, conj(a₀),
                                     conj(a₁))
-        if info.converged > 0 && reverse_info.converged == 0 && alg_rrule.verbosity >= 0
+        if info.converged > 0 && reverse_info.converged == 0 && alg_primal.verbosity >= 1
             @warn "`linsolve` cotangent problem did not converge, whereas the primal linear problem did: normres = $(reverse_info.normres)"
         end
         x∂b = inner(x, ∂b)

ext/KrylovKitChainRulesCoreExt/svdsolve.jl

@@ -112,7 +112,7 @@ function compute_svdsolve_pullback_data(Δvals, Δlvecs, Δrvecs, vals, lvecs, r
         udΔu = inner(u, Δu)
         vdΔv = inner(v, Δv)
         if (udΔu isa Complex) || (vdΔv isa Complex)
-            if alg_rrule.verbosity >= 0
+            if alg_primal.verbosity >= 1
                 gauge = abs(imag(udΔu + vdΔv))
                 gauge > alg_primal.tol &&
                     @warn "`svdsolve` cotangents for singular vectors $i are sensitive to gauge choice: (|gauge| = $gauge)"
@@ -131,7 +131,7 @@ function compute_svdsolve_pullback_data(Δvals, Δlvecs, Δrvecs, vals, lvecs, r
                 return (x′, y′)
             end
         end
-        if info.converged >= i && reverse_info.converged == 0 && alg_rrule.verbosity >= 0
+        if info.converged >= i && reverse_info.converged == 0 && alg_primal.verbosity >= 0
             @warn "`svdsolve` cotangent linear problem ($i) did not converge, whereas the primal eigenvalue problem did: normres = $(reverse_info.normres)"
         end
         x = VectorInterface.add!!(x, u, Δs / 2)
@@ -162,7 +162,7 @@ function compute_svdsolve_pullback_data(Δvals, Δlvecs, Δrvecs, vals, lvecs, r
     aVdΔV = rmul!(VdΔV - VdΔV', 1 / 2)
 
     tol = alg_primal.tol
-    if alg_rrule.verbosity >= 0
+    if alg_primal.verbosity >= 1
         mask = abs.(vals' .- vals) .< tol
         gaugepart = view(aUdΔU, mask) + view(aVdΔV, mask)
         gauge = norm(gaugepart, Inf)
@@ -227,7 +227,7 @@ function compute_svdsolve_pullback_data(Δvals, Δlvecs, Δrvecs, vals, lvecs, r
             return (x′, y′, vals .* z)
         end
     end
-    if info.converged >= n && reverse_info.converged < n && alg_rrule.verbosity >= 0
+    if info.converged >= n && reverse_info.converged < n && alg_primal.verbosity >= 1
         @warn "`svdsolve` cotangent problem did not converge, whereas the primal singular value problem did"
     end
 
@@ -236,13 +236,13 @@ function compute_svdsolve_pullback_data(Δvals, Δlvecs, Δrvecs, vals, lvecs, r
     for i in 1:n
         x, y, z = Ws[i]
         _, ic = findmax(abs, z)
-        if ic != i
+        if ic != i && alg_primal.verbosity >= 1
             @warn "`svdsolve` cotangent linear problem ($ic) returns unexpected result"
         end
         factor = 1 / z[ic]
         z[ic] = zero(z[ic])
         error = max(norm(z, Inf), abs(rvals[i] - vals[ic]))
-        if error > 5 * tol && alg_rrule.verbosity >= 0
+        if error > 10 * tol && alg_primal.verbosity >= 1
             @warn "`svdsolve` cotangent linear problem ($ic) returns unexpected result: error = $error vs tol = $tol"
         end
         xs[ic] = VectorInterface.add!!(xs[ic], x, -factor)

test/ad/degenerateeigsolve.jl

@@ -94,8 +94,8 @@ end
 
     tol = 2 * N^2 * eps(real(T))
     alg = Arnoldi(; tol=tol, krylovdim=2n)
-    alg_rrule1 = Arnoldi(; tol=tol, krylovdim=2n, verbosity=-1)
-    alg_rrule2 = GMRES(; tol=tol, krylovdim=2n, verbosity=-1)
+    alg_rrule1 = Arnoldi(; tol=tol, krylovdim=2n)
+    alg_rrule2 = GMRES(; tol=tol, krylovdim=2n)
     mat_example1, mat_example_fun1, mat_example_fd, Avec, Bvec, Cvec, xvec, vals, vecs = build_mat_example(A,
                                                                                                            B,
                                                                                                            C,

test/ad/eigsolve.jl

@@ -205,8 +205,8 @@ end
     condA = cond(A)
     tol = n * condA * (T <: Real ? eps(T) : 4 * eps(real(T)))
     alg = Arnoldi(; tol=tol, krylovdim=n)
-    alg_rrule1 = Arnoldi(; tol=tol, krylovdim=2n, verbosity=-1)
-    alg_rrule2 = GMRES(; tol=tol, krylovdim=n + 1, verbosity=-1)
+    alg_rrule1 = Arnoldi(; tol=tol, krylovdim=2n)
+    alg_rrule2 = GMRES(; tol=tol, krylovdim=n + 1)
     config = Zygote.ZygoteRuleConfig()
     @testset for which in whichlist
         for alg_rrule in (alg_rrule1, alg_rrule2)
@@ -269,11 +269,13 @@ end
 
     if T <: Complex
         @testset "test warnings and info" begin
-            alg_rrule = Arnoldi(; tol=tol, krylovdim=n, verbosity=-1)
+            alg = Arnoldi(; tol=tol, krylovdim=n, verbosity=0)
+            alg_rrule = Arnoldi(; tol=tol, krylovdim=n, verbosity=0)
             (vals, vecs, info), pb = ChainRulesCore.rrule(config, eigsolve, A, x, howmany,
                                                           :LR, alg; alg_rrule=alg_rrule)
             @test_logs pb((ZeroTangent(), im .* vecs[1:2] .+ vecs[2:-1:1], NoTangent()))
 
+            alg = Arnoldi(; tol=tol, krylovdim=n, verbosity=1)
             alg_rrule = Arnoldi(; tol=tol, krylovdim=n, verbosity=0)
             (vals, vecs, info), pb = ChainRulesCore.rrule(config, eigsolve, A, x, howmany,
                                                           :LR, alg; alg_rrule=alg_rrule)
@@ -282,6 +284,7 @@ end
             pbs = @test_logs pb((ZeroTangent(), vecs[1:2], NoTangent()))
             @test norm(unthunk(pbs[1]), Inf) < condA * sqrt(eps(real(T)))
 
+            alg = Arnoldi(; tol=tol, krylovdim=n, verbosity=1)
             alg_rrule = Arnoldi(; tol=tol, krylovdim=n, verbosity=1)
             (vals, vecs, info), pb = ChainRulesCore.rrule(config, eigsolve, A, x, howmany,
                                                           :LR, alg; alg_rrule=alg_rrule)
@@ -290,11 +293,13 @@ end
             pbs = @test_logs (:info,) pb((ZeroTangent(), vecs[1:2], NoTangent()))
             @test norm(unthunk(pbs[1]), Inf) < condA * sqrt(eps(real(T)))
 
-            alg_rrule = GMRES(; tol=tol, krylovdim=n, verbosity=-1)
+            alg = Arnoldi(; tol=tol, krylovdim=n, verbosity=0)
+            alg_rrule = GMRES(; tol=tol, krylovdim=n, verbosity=0)
             (vals, vecs, info), pb = ChainRulesCore.rrule(config, eigsolve, A, x, howmany,
                                                           :LR, alg; alg_rrule=alg_rrule)
             @test_logs pb((ZeroTangent(), im .* vecs[1:2] .+ vecs[2:-1:1], NoTangent()))
 
+            alg = Arnoldi(; tol=tol, krylovdim=n, verbosity=1)
             alg_rrule = GMRES(; tol=tol, krylovdim=n, verbosity=0)
             (vals, vecs, info), pb = ChainRulesCore.rrule(config, eigsolve, A, x, howmany,
                                                           :LR, alg; alg_rrule=alg_rrule)
@@ -305,6 +310,7 @@ end
             pbs = @test_logs pb((ZeroTangent(), vecs[1:2], NoTangent()))
             @test norm(unthunk(pbs[1]), Inf) < condA * sqrt(eps(real(T)))
 
+            alg = Arnoldi(; tol=tol, krylovdim=n, verbosity=1)
             alg_rrule = GMRES(; tol=tol, krylovdim=n, verbosity=1)
             (vals, vecs, info), pb = ChainRulesCore.rrule(config, eigsolve, A, x, howmany,
                                                           :LR, alg; alg_rrule=alg_rrule)

test/ad/svdsolve.jl

@@ -154,8 +154,8 @@ end
     howmany = 3
     tol = 3 * n * condA * (T <: Real ? eps(T) : 4 * eps(real(T)))
     alg = GKL(; krylovdim=2n, tol=tol)
-    alg_rrule1 = Arnoldi(; tol=tol, krylovdim=4n, verbosity=-1)
-    alg_rrule2 = GMRES(; tol=tol, krylovdim=3n, verbosity=-1)
+    alg_rrule1 = Arnoldi(; tol=tol, krylovdim=4n)
+    alg_rrule2 = GMRES(; tol=tol, krylovdim=3n)
     config = Zygote.ZygoteRuleConfig()
     for alg_rrule in (alg_rrule1, alg_rrule2)
         # unfortunately, rrule does not seem type stable for function arguments, because the
@@ -219,13 +219,15 @@ end
     end
     if T <: Complex
         @testset "test warnings and info" begin
-            alg_rrule = Arnoldi(; tol=tol, krylovdim=4n, verbosity=-1)
+            alg = GKL(; krylovdim=2n, tol=tol, verbosity=0)
+            alg_rrule = Arnoldi(; tol=tol, krylovdim=4n, verbosity=0)
             (vals, lvecs, rvecs, info), pb = ChainRulesCore.rrule(config, svdsolve, A, x,
                                                                   howmany, :LR, alg;
                                                                   alg_rrule=alg_rrule)
             @test_logs pb((ZeroTangent(), im .* lvecs[1:2] .+ lvecs[2:-1:1], ZeroTangent(),
                            NoTangent()))
 
+            alg = GKL(; krylovdim=2n, tol=tol, verbosity=1)
             alg_rrule = Arnoldi(; tol=tol, krylovdim=4n, verbosity=0)
             (vals, lvecs, rvecs, info), pb = ChainRulesCore.rrule(config, svdsolve, A, x,
                                                                   howmany, :LR, alg;
@@ -249,6 +251,7 @@ end
                            (1 - im) .* rvecs[1:2] + rvecs[2:-1:1],
                            NoTangent()))
 
+            alg = GKL(; krylovdim=2n, tol=tol, verbosity=1)
             alg_rrule = Arnoldi(; tol=tol, krylovdim=4n, verbosity=1)
             (vals, lvecs, rvecs, info), pb = ChainRulesCore.rrule(config, svdsolve, A, x,
                                                                   howmany, :LR, alg;
@@ -272,13 +275,15 @@ end
                                     (1 - im) .* rvecs[1:2] + rvecs[2:-1:1],
                                     NoTangent()))
 
-            alg_rrule = GMRES(; tol=tol, krylovdim=3n, verbosity=-1)
+            alg = GKL(; krylovdim=2n, tol=tol, verbosity=0)
+            alg_rrule = GMRES(; tol=tol, krylovdim=3n, verbosity=0)
             (vals, lvecs, rvecs, info), pb = ChainRulesCore.rrule(config, svdsolve, A, x,
                                                                   howmany, :LR, alg;
                                                                   alg_rrule=alg_rrule)
             @test_logs pb((ZeroTangent(), im .* lvecs[1:2] .+ lvecs[2:-1:1], ZeroTangent(),
                            NoTangent()))
 
+            alg = GKL(; krylovdim=2n, tol=tol, verbosity=1)
             alg_rrule = GMRES(; tol=tol, krylovdim=3n, verbosity=0)
             (vals, lvecs, rvecs, info), pb = ChainRulesCore.rrule(config, svdsolve, A, x,
                                                                   howmany, :LR, alg;
@@ -305,6 +310,7 @@ end
                            (1 - im) .* rvecs[1:2] + rvecs[2:-1:1],
                            NoTangent()))
 
+            alg = GKL(; krylovdim=2n, tol=tol, verbosity=1)
             alg_rrule = GMRES(; tol=tol, krylovdim=3n, verbosity=1)
             (vals, lvecs, rvecs, info), pb = ChainRulesCore.rrule(config, svdsolve, A, x,
                                                                   howmany, :LR, alg;

test/testsetup.jl

@@ -30,83 +30,20 @@ function buildrealmap(A, B)
     return x -> A * x + B * conj(x)
 end
 
-# Minimal vector type
-# -------------------
-"""
-    MinimalVec{T<:Number,IP}
-
-Minimal interface for a vector. Can support either in-place assignments or not, depending on
-`IP=true` or `IP=false`.
-"""
-struct MinimalVec{IP,V<:AbstractVector}
-    vec::V
-    function MinimalVec{IP}(vec::V) where {IP,V}
-        return new{IP,V}(vec)
-    end
-end
-const InplaceVec{V} = MinimalVec{true,V}
-const OutplaceVec{V} = MinimalVec{false,V}
-
-isinplace(::Type{MinimalVec{IP,V}}) where {V,IP} = IP
-isinplace(v::MinimalVec) = isinplace(typeof(v))
-
-VI.scalartype(::Type{<:MinimalVec{IP,V}}) where {IP,V} = scalartype(V)
-
-function VI.zerovector(v::MinimalVec, S::Type{<:Number})
-    return MinimalVec{isinplace(v)}(zerovector(v.vec, S))
-end
-function VI.zerovector!(v::InplaceVec{V}) where {V}
-    zerovector!(v.vec)
-    return v
-end
-VI.zerovector!!(v::MinimalVec) = isinplace(v) ? zerovector!(v) : zerovector(v)
-
-function VI.scale(v::MinimalVec, α::Number)
-    return MinimalVec{isinplace(v)}(scale(v.vec, α))
-end
-function VI.scale!(v::InplaceVec{V}, α::Number) where {V}
-    scale!(v.vec, α)
-    return v
-end
-function VI.scale!!(v::MinimalVec, α::Number)
-    return isinplace(v) ? scale!(v, α) : scale(v, α)
-end
-function VI.scale!(w::InplaceVec{V}, v::InplaceVec{W}, α::Number) where {V,W}
-    scale!(w.vec, v.vec, α)
-    return w
-end
-function VI.scale!!(w::MinimalVec, v::MinimalVec, α::Number)
-    isinplace(w) && return scale!(w, v, α)
-    return MinimalVec{false}(scale!!(copy(w.vec), v.vec, α))
-end
-
-function VI.add(y::MinimalVec, x::MinimalVec, α::Number, β::Number)
-    return MinimalVec{isinplace(y)}(add(y.vec, x.vec, α, β))
-end
-function VI.add!(y::InplaceVec{W}, x::InplaceVec{V}, α::Number, β::Number) where {W,V}
-    add!(y.vec, x.vec, α, β)
-    return y
-end
-function VI.add!!(y::MinimalVec, x::MinimalVec, α::Number, β::Number)
-    return isinplace(y) ? add!(y, x, α, β) : add(y, x, α, β)
-end
-
-VI.inner(x::MinimalVec, y::MinimalVec) = inner(x.vec, y.vec)
-VI.norm(x::MinimalVec) = LinearAlgebra.norm(x.vec)
-
 # Wrappers
 # --------
+using VectorInterface: MinimalSVec, MinimalMVec, MinimalVec
 # dispatch on val is necessary for type stability
 
 function wrapvec(v, ::Val{mode}) where {mode}
     return mode === :vector ? v :
-           mode === :inplace ? MinimalVec{true}(v) :
-           mode === :outplace ? MinimalVec{false}(v) :
-           mode === :mixed ? MinimalVec{false}(v) :
+           mode === :inplace ? MinimalMVec(v) :
+           mode === :outplace ? MinimalSVec(v) :
+           mode === :mixed ? MinimalSVec(v) :
            throw(ArgumentError("invalid mode ($mode)"))
 end
 function wrapvec2(v, ::Val{mode}) where {mode}
-    return mode === :mixed ? MinimalVec{true}(v) : wrapvec(v, mode)
+    return mode === :mixed ? MinimalMVec(v) : wrapvec(v, mode)
 end
 
 unwrapvec(v::MinimalVec) = v.vec