SVD pullback of rank-deficient matrix

[lkdvos]

I encountered a DimensionMismatch error during the backward pass of an optimization loop involving SVD truncation on $U(1)$ symmetric tensors. I think this issue occurs when the distribution of bond dimensions across different charge sectors changes between the forward pass and the backward pullback. As I don't want to fixed the sectors, I really need this function maybe.

ERROR: LoadError: DimensionMismatch: 
Stacktrace:
  [1] svd_pullback!(ΔA::Base.ReshapedArray{Float64, 2, SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}}, true}, Tuple{}}, A::Base.ReshapedArray{Float64, 2, SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}}, true}, Tuple{}}, USVᴴ::Tuple{Base.ReshapedArray{Float64, 2, SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}}, true}, Tuple{}}, Diagonal{Float64, SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}}, true}}, Base.ReshapedArray{Float64, 2, SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}}, true}, Tuple{}}}, ΔUSVᴴ::Vector{Any}, ind::BitVector; tol::Float64, rank_atol::Float64, degeneracy_atol::Float64, gauge_atol::Float64)
    @ MatrixAlgebraKit ~/.julia/packages/MatrixAlgebraKit/xv5dv/src/pullbacks/svd.jl:51
  [2] svd_pullback!(ΔA::Base.ReshapedArray{Float64, 2, SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}}, true}, Tuple{}}, A::Base.ReshapedArray{Float64, 2, SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}}, true}, Tuple{}}, USVᴴ::Tuple{Base.ReshapedArray{Float64, 2, SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}}, true}, Tuple{}}, Diagonal{Float64, SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}}, true}}, Base.ReshapedArray{Float64, 2, SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}}, true}, Tuple{}}}, ΔUSVᴴ::Vector{Any}, ind::BitVector)
    @ MatrixAlgebraKit ~/.julia/packages/MatrixAlgebraKit/xv5dv/src/pullbacks/svd.jl:24

Originally posted by @Ryo-wtnb11 in #101

I investigated this issue in a bit more detail and can distill the problem to the following MWE:

using Zygote, MatrixAlgebraKit

# generate a rank-deficient matrix
A = zeros(3, 3)
rand!(@view(A[1:2, 1:2]))

# dummy function to derive
function f(A)
    U, S, Vh = svd_compact(A)
    return tr(U * S * Vh)
end

Zygote.gradient(f, A)
ERROR: DimensionMismatch: 
Stacktrace:
  [1] svd_pullback!(ΔA::Matrix{…}, A::Matrix{…}, USVᴴ::Tuple{…}, ΔUSVᴴ::Vector{…}, ind::Colon; tol::Float64, rank_atol::Float64, degeneracy_atol::Float64, gauge_atol::Float64)
    @ MatrixAlgebraKit ~/.julia/packages/MatrixAlgebraKit/xv5dv/src/pullbacks/svd.jl:51
  [2] svd_pullback!
    @ ~/.julia/packages/MatrixAlgebraKit/xv5dv/src/pullbacks/svd.jl:24 [inlined]
  [3] svd_pullback!(ΔA::Matrix{…}, A::Matrix{…}, USVᴴ::Tuple{…}, ΔUSVᴴ::Vector{…})
    @ MatrixAlgebraKit ~/.julia/packages/MatrixAlgebraKit/xv5dv/src/pullbacks/svd.jl:24

---

I think I have an idea of what is going on, I am pretty sure I can bypass the error but I'm not entirely sure everything will still be correct.
The main point is that we are currently checking that we are not keeping more singular values than the numerical rank of the matri.
This happens when indU or indV has more entries than the rank, throwing a DimensionMismatch for that (which is clearly not the most informative message either...).

In principle I might be able to alter the implementation to ignore that, but I'm not entirely sure if that is mathematically allowed, since this is basically one large degenerate subspace (nullspace) and I don't know if this would be correctly taking that into account.

@Jutho, I don't know if you have more insight into this?

[Jutho]

Yes, this is clearly a bug. Revisiting the svd_pullback! was anyway on my todo list, as the current implementation does also not work as a pullback for svd_full in the rectangular case, even if A is full rank.

However, if the rank of A is smaller than minmn - 1, i.e. if at least two singular values are zero, there is a true degeneracy, which still causes some issues in the derivatives of the singular vectors, for which I don't really have a satisfactory solution (yet). This is similar to the (hermitian) eigenvector case.

[Ryo-wtnb11]

I've tried using trunctol() to discard zero singular values. While it works in most cases, the same issue still persists in certain situations. Is there a more fundamental problem at play here? Does svd_pullback! need a specific way to treat the null space in rectangular matrices?