Error after unconverged KrylovKit.GMRES linear solve

[pbrehmer]

It seems that on the newest version KrylovKit something broke: If we compute the fixed-point gradient fpgrad using KrylovKit.GMRES and it doesn't converge, I get the following conversion error:

ERROR: MethodError: Cannot `convert` an object of type
  CTMRGEnv{TensorMap{ComplexSpace,1,1,Trivial,Array{Complex{Float64},2},Nothing,Nothing},TensorMap{ComplexSpace,1,1,Trivial,Array{Complex{Float64},2},Nothing,Nothing}} to an object of type
  CTMRGEnv{TensorMap{ComplexSpace,1,1,Trivial,Array{Complex{Float64},2},Nothing,Nothing},TensorMap{ComplexSpace,3,1,Trivial,Array{Complex{Float64},2},Nothing,Nothing}}

Closest candidates are:
  convert(::Type{T}, ::T) where T
   @ Base Base.jl:84
  (::Type{CTMRGEnv{C, T}} where {C, T})(::Any, ::Any)
   @ PEPSKit ~/repos/PEPSKit.jl/src/environments/ctmrgenv.jl:7

Stacktrace:
  [1] setindex!(A::Vector{CTMRGEnv{…}}, x::CTMRGEnv{TrivialTensorMap{…}, TrivialTensorMap{…}}, i1::Int64)
    @ Base ./array.jl:1021
  [2] setindex!(b::KrylovKit.OrthonormalBasis{CTMRGEnv{…}}, i::CTMRGEnv{TrivialTensorMap{…}, TrivialTensorMap{…}}, q::Int64)
    @ KrylovKit ~/.julia/packages/KrylovKit/wPepJ/src/orthonormal.jl:42
  [3] _rmul!(b::KrylovKit.OrthonormalBasis{CTMRGEnv{…}}, G::LinearAlgebra.Givens{ComplexF64})
    @ KrylovKit ~/.julia/packages/KrylovKit/wPepJ/src/dense/givens.jl:34
  [4] rmul!(b::KrylovKit.OrthonormalBasis{CTMRGEnv{…}}, G::LinearAlgebra.Givens{ComplexF64})
    @ KrylovKit ~/.julia/packages/KrylovKit/wPepJ/src/dense/givens.jl:16
  [5] linsolve(operator::Function, b::CTMRGEnv{…}, x₀::CTMRGEnv{…}, alg::GMRES{…}, a₀::Int64, a₁::Int64)
    @ KrylovKit ~/.julia/packages/KrylovKit/wPepJ/src/linsolve/gmres.jl:114
  [6] fpgrad(∂F∂x::CTMRGEnv{…}, ∂f∂x::Function, ∂f∂A::PEPSKit.var"#∂f∂A#174"{…}, y₀::CTMRGEnv{…}, alg::GMRES{…})
    @ PEPSKit ~/repos/PEPSKit.jl/src/algorithms/peps_opt.jl:201
  [7] (::PEPSKit.var"#leading_boundary_pullback#172"{…})(Δenvs′::CTMRGEnv{…})
    @ PEPSKit ~/repos/PEPSKit.jl/src/algorithms/peps_opt.jl:128
  [8] #28
    @ ~/repos/PEPSKit.jl/src/utility/hook_pullback.jl:34 [inlined]
  [9] ZBack
    @ ~/.julia/packages/Zygote/nsBv0/src/compiler/chainrules.jl:211 [inlined]
 [10] (::Zygote.var"#kw_zpullback#53"{PEPSKit.var"#28#29"{…}})(dy::CTMRGEnv{TrivialTensorMap{…}, TrivialTensorMap{…}})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/chainrules.jl:237
 [11] #2
    @ ~/repos/PEPSKit.jl/test/ctmrg/gradients.jl:63 [inlined]
 [12] (::Zygote.Pullback{Tuple{…}, Tuple{…}})(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0
 [13] (::Zygote.var"#75#76"{Zygote.Pullback{Tuple{…}, Tuple{…}}})(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface.jl:91
 [14] withgradient(f::Function, args::InfinitePEPS{TrivialTensorMap{ComplexSpace, 1, 4, Matrix{ComplexF64}}})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface.jl:213
 [15] compute_grad(peps::InfinitePEPS{…}, envs::CTMRGEnv{…}, boundary_alg::CTMRG, alg_rrule::GMRES{…})
    @ Main ~/repos/PEPSKit.jl/test/ctmrg/gradients.jl:62
 [16] top-level scope
    @ ~/repos/PEPSKit.jl/test/ctmrg/gradients.jl:70

It turns out that this only happens on the newest version v0.7.1, whereas on the previous version v0.6.1 this was not a problem. For comparison, I currently use the following package versions:

(PEPSKit) pkg> status
Project PEPSKit v0.1.0
Status `~/repos/PEPSKit.jl/Project.toml`
  [7d9f7c33] Accessors v0.1.36
  [d360d2e6] ChainRulesCore v1.24.0
  [34da2185] Compat v4.15.0
  [0b1a1467] KrylovKit v0.7.1
  [bb1c41ca] MPSKit v0.10.2 `https://github.com/maartenvd/MPSKit.jl.git#master`
  [77e91f04] OptimKit v0.3.1
  [07d1fe3e] TensorKit v0.12.4
  [409d34a3] VectorInterface v0.4.5
  [e88e6eb3] Zygote v0.6.70
  [37e2e46d] LinearAlgebra
  [de0858da] Printf
  [9a3f8284] Random
  [10745b16] Statistics v1.10.0

Note that I added the master branch of MPSKit to resolve the KrylovKit dependency conflict. Here's a MWE that should reproduce the error where I intentionally prevent the linear solver from converging:

using Random
using PEPSKit
using TensorKit
using Zygote
using KrylovKit

Random.seed!(42039482038)
χbond = 2
χenv = 8
boundary_alg = CTMRG(;
    trscheme=truncdim(χenv), tol=1e-10, miniter=20, maxiter=100, verbosity=1
)
psi = InfinitePEPS(2, χbond)
env = leading_boundary(CTMRGEnv(psi; Venv=ℂ^χenv), psi, boundary_alg)

function compute_grad(peps, envs, boundary_alg, alg_rrule)
    E, g = Zygote.withgradient(peps) do psi
        envs2 = PEPSKit.hook_pullback(leading_boundary, envs, psi, boundary_alg; alg_rrule)
        return costfun(psi, envs2, H)
    end
    return only(g)
end

g = compute_grad(psi, env, boundary_alg, KrylovKit.GMRES(; tol=1e-16, maxiter=2, verbosity=3))

Maybe @lkdvos knows what's going on here?

[lkdvos]

I am not sure why this only ever happened when it does not converge, but this fix seems to at least solve the minimal example you added (thanks for that by the way, that made solving this so much easier!)

[pbrehmer]

Thanks for the super quick fix!