Correct adjoint for truncated SVD

src/utility/svd.jl

@@ -26,18 +26,19 @@
 # Wrapper around Krylov Kit's GKL iterative SVD solver
 @kwdef struct IterSVD
     alg::KrylovKit.GKL = KrylovKit.GKL(; tol=1e-14, krylovdim=25)
+    fallback_threshold::Float64 = Inf
     lorentz_broad::Float64 = 0.0
     alg_rrule::Union{GMRES,BiCGStab,Arnoldi} = GMRES(; tol=1e-14)
 end
 
 # Compute SVD data block-wise using KrylovKit algorithm
 function TensorKit._tsvd!(
-    t, alg::IterSVD, trunc::Union{NoTruncation,TruncationSpace}, p::Real=2
+    t, alg::Union{IterSVD}, trunc::Union{NoTruncation,TruncationSpace}, p::Real=2
 )
     # early return
     if isempty(blocksectors(t))
         truncerr = zero(real(scalartype(t)))
         return _empty_svdtensors(t)..., truncerr
     end
 
     Udata, Σdata, Vdata, dims = _compute_svddata!(t, alg, trunc)
@@ -55,22 +56,30 @@
     Udata = SectorDict{I,A}()
     Vdata = SectorDict{I,A}()
     dims = SectorDict{I,Int}()
     local Sdata
     for (c, b) in blocks(t)
         x₀ = randn(eltype(b), size(b, 1))
         howmany = trunc isa NoTruncation ? minimum(size(b)) : blockdim(trunc.space, c)
-        S, lvecs, rvecs, info = KrylovKit.svdsolve(b, x₀, howmany, :LR, alg.alg)
-        if info.converged < howmany  # Fall back to dense SVD if not properly converged
+
+        if howmany / minimum(size(b)) > alg.fallback_threshold  # Use dense SVD for small blocks
             U, S, V = TensorKit.MatrixAlgebra.svd!(b, TensorKit.SVD())
-            Udata[c] = U
-            Vdata[c] = V
-        else  # Slice in case more values were converged than requested
-            Udata[c] = stack(view(lvecs, 1:howmany))
-            Vdata[c] = stack(conj, view(rvecs, 1:howmany); dims=1)
-            S = @view S[1:howmany]
+            Udata[c] = @view U[:, 1:howmany]
+            Vdata[c] = @view V[1:howmany, :]
+        else
+            S, lvecs, rvecs, info = KrylovKit.svdsolve(b, x₀, howmany, :LR, alg.alg)
+            if info.converged < howmany  # Fall back to dense SVD if not properly converged
+                U, S, V = TensorKit.MatrixAlgebra.svd!(b, TensorKit.SVD())
+                Udata[c] = @view U[:, 1:howmany]
+                Vdata[c] = @view V[1:howmany, :]
+            else  # Slice in case more values were converged than requested
+                Udata[c] = stack(view(lvecs, 1:howmany))
+                Vdata[c] = stack(conj, view(rvecs, 1:howmany); dims=1)
+            end
         end
+
+        S = @view S[1:howmany]
         if @isdefined Sdata # cannot easily infer the type of Σ, so use this construction
             Sdata[c] = S
         else
             Sdata = SectorDict(c => S)
         end
@@ -103,8 +112,8 @@
             minimal_info = KrylovKit.ConvergenceInfo(length(Sdc), nothing, nothing, -1, -1)  # Just supply converged to SVD pullback
 
             if ΔUc isa AbstractZero && ΔVc isa AbstractZero  # Handle ZeroTangent singular vectors
                 Δlvecs = fill(ZeroTangent(), n_vals)
                 Δrvecs = fill(ZeroTangent(), n_vals)
             else
                 Δlvecs = Vector{Vector{scalartype(t)}}(eachcol(ΔUc))
                 Δrvecs = Vector{Vector{scalartype(t)}}(eachcol(ΔVc'))
@@ -130,7 +139,7 @@
         end
         return NoTangent(), Δt, NoTangent()
     end
     function tsvd_itersvd_pullback(::Tuple{ZeroTangent,ZeroTangent,ZeroTangent})
         return NoTangent(), ZeroTangent(), NoTangent()
     end
 
@@ -169,7 +178,7 @@
         end
         return NoTangent(), Δt, NoTangent()
     end
     function tsvd_oldsvd_pullback(::Tuple{ZeroTangent,ZeroTangent,ZeroTangent})
         return NoTangent(), ZeroTangent(), NoTangent()
     end
 
@@ -235,7 +244,7 @@
 end
 
 # Lorentzian broadening for SVD adjoint F-singularities
 function _lorentz_broaden(x::Real, ε=1e-12)
     x′ = 1 / x
     return x′ / (x′^2 + ε)
 end