Updates for eig and lq/qr diag pullbacks

Author: kshyatt

src/pullbacks/eig.jl

@@ -1,12 +1,3 @@
-function check_eig_cotangents(D, VᴴΔV; degeneracy_atol::Real = default_pullback_rank_atol(D), gauge_atol::Real = default_pullback_gauge_atol(VᴴΔV))
-    mask = abs.(transpose(D) .- D) .< degeneracy_atol
-    # not GPU friendly...
-    Δgauge = norm(view(VᴴΔV, mask))
-    Δgauge ≤ gauge_atol ||
-        @warn "`eig` cotangents sensitive to gauge choice: (|Δgauge| = $Δgauge)"
-    return
-end
-
 """
     eig_pullback!(
         ΔA::AbstractMatrix, A, DV, ΔDV, [ind];
@@ -49,7 +40,11 @@ function eig_pullback!(
         indV = axes(V, 2)[ind]
         length(indV) == pV || throw(DimensionMismatch())
         mul!(view(VᴴΔV, :, indV), V', ΔV)
-        check_eig_cotangents(D, VᴴΔV; degeneracy_atol, gauge_atol)
+
+        mask = abs.(transpose(D) .- D) .< degeneracy_atol
+        Δgauge = norm(view(VᴴΔV, mask), Inf)
+        Δgauge ≤ gauge_atol ||
+            @warn "`eig` cotangents sensitive to gauge choice: (|Δgauge| = $Δgauge)"
 
         VᴴΔV .*= conj.(inv_safe.(transpose(D) .- D, degeneracy_atol))
 
@@ -83,16 +78,6 @@ function eig_pullback!(
     end
     return ΔA
 end
-function eig_pullback!(
-        ΔA::Diagonal, A, DV, ΔDV, ind = Colon();
-        degeneracy_atol::Real = default_pullback_rank_atol(DV[1]),
-        gauge_atol::Real = default_pullback_gauge_atol(ΔDV[2])
-    )
-    ΔA_full = zero!(similar(ΔA, size(ΔA)))
-    ΔA_full = eig_pullback!(ΔA_full, A, DV, ΔDV, ind; degeneracy_atol, gauge_atol)
-    diagview(ΔA) .+= diagview(ΔA_full)
-    return ΔA
-end
 
 """
     eig_trunc_pullback!(
@@ -134,7 +119,10 @@ function eig_trunc_pullback!(
     if !iszerotangent(ΔV)
         (n, p) == size(ΔV) || throw(DimensionMismatch())
         VᴴΔV = V' * ΔV
-        check_eig_cotangents(D, VᴴΔV; degeneracy_atol, gauge_atol)
+        mask = abs.(transpose(D) .- D) .< degeneracy_atol
+        Δgauge = norm(view(VᴴΔV, mask), Inf)
+        Δgauge ≤ gauge_atol ||
+            @warn "`eig` cotangents sensitive to gauge choice: (|Δgauge| = $Δgauge)"
 
         ΔVperp = ΔV - V * inv(G) * VᴴΔV
         VᴴΔV .*= conj.(inv_safe.(transpose(D) .- D, degeneracy_atol))
@@ -163,16 +151,6 @@ function eig_trunc_pullback!(
     end
     return ΔA
 end
-function eig_trunc_pullback!(
-        ΔA::Diagonal, A, DV, ΔDV;
-        degeneracy_atol::Real = default_pullback_rank_atol(DV[1]),
-        gauge_atol::Real = default_pullback_gauge_atol(ΔDV[2])
-    )
-    ΔA_full = zero!(similar(ΔA, size(ΔA)))
-    ΔA_full = eig_trunc_pullback!(ΔA_full, A, DV, ΔDV; degeneracy_atol, gauge_atol)
-    diagview(ΔA) .+= diagview(ΔA_full)
-    return ΔA
-end
 
 """
     eig_vals_pullback!(
@@ -193,6 +171,29 @@ function eig_vals_pullback!(
         ΔA, A, DV, ΔD, ind = Colon();
         degeneracy_atol::Real = default_pullback_rank_atol(DV[1]),
     )
+
     ΔDV = (diagonal(ΔD), nothing)
     return eig_pullback!(ΔA, A, DV, ΔDV, ind; degeneracy_atol)
 end
+
+function eig_pullback!(
+        ΔA::Diagonal, A, DV, ΔDV, ind = Colon();
+        degeneracy_atol::Real = default_pullback_rank_atol(DV[1]),
+        gauge_atol::Real = default_pullback_gauge_atol(ΔDV[2])
+    )
+    ΔA_full = zero!(similar(ΔA, size(ΔA)))
+    eig_pullback!(ΔA_full, A, DV, ΔDV, ind; degeneracy_atol, gauge_atol)
+    diagview(ΔA) .+= diagview(ΔA_full)
+    return ΔA
+end
+
+function eig_trunc_pullback!(
+        ΔA::Diagonal, A, DV, ΔDV;
+        degeneracy_atol::Real = default_pullback_rank_atol(DV[1]),
+        gauge_atol::Real = default_pullback_gauge_atol(ΔDV[2])
+    )
+    ΔA_full = zero!(similar(ΔA, size(ΔA)))
+    eig_trunc_pullback!(ΔA_full, A, DV, ΔDV; degeneracy_atol, gauge_atol)
+    diagview(ΔA) .+= diagview(ΔA_full)
+    return ΔA
+end

src/pullbacks/lq.jl

@@ -124,6 +124,16 @@ function lq_pullback!(
     ΔA1 .+= ΔQ̃
     return ΔA
 end
+function lq_pullback!(
+        ΔA::Diagonal, A, LQ, ΔLQ;
+        rank_atol::Real = default_pullback_rank_atol(LQ[1]),
+        gauge_atol::Real = default_pullback_gauge_atol(ΔLQ[2])
+    )
+    ΔA_full = zero!(similar(ΔA, size(ΔA)))
+    ΔA_full = lq_pullback!(ΔA_full, A, LQ, ΔLQ; rank_atol, gauge_atol)
+    diagview(ΔA) .+= diagview(ΔA_full)
+    return ΔA
+end
 
 function check_lq_null_cotangents(Nᴴ, ΔNᴴ; gauge_atol::Real = default_pullback_gauge_atol(ΔNᴴ))
     aNᴴΔN = project_antihermitian!(Nᴴ * ΔNᴴ')

src/pullbacks/polar.jl

@@ -27,6 +27,12 @@ function left_polar_pullback!(ΔA::AbstractMatrix, A, WP, ΔWP; kwargs...)
     end
     return ΔA
 end
+function left_polar_pullback!(ΔA::Diagonal, A, WP, ΔWP; kwargs...)
+    ΔA_full = zero!(similar(ΔA, size(ΔA)))
+    ΔA_full = left_polar_pullback!(ΔA_full, A, WP, ΔWP; kwargs...)
+    diagview(ΔA) .+= diagview(ΔA_full)
+    return ΔA
+end
 
 """
     right_polar_pullback!(ΔA, A, PWᴴ, ΔPWᴴ)
@@ -57,3 +63,9 @@ function right_polar_pullback!(ΔA::AbstractMatrix, A, PWᴴ, ΔPWᴴ; kwargs...
     end
     return ΔA
 end
+function right_polar_pullback!(ΔA::Diagonal, A, PWᴴ, ΔPWᴴ; kwargs...)
+    ΔA_full = zero!(similar(ΔA, size(ΔA)))
+    ΔA_full = right_polar_pullback!(ΔA_full, A, PWᴴ, ΔPWᴴ; kwargs...)
+    diagview(ΔA) .+= diagview(ΔA_full)
+    return ΔA
+end

src/pullbacks/qr.jl

@@ -10,7 +10,7 @@ function check_qr_cotangents(Q, R, ΔQ, ΔR, minmn::Int, p::Int; gauge_atol::Rea
             Δgauge = max(Δgauge, norm(ΔQ2, Inf))
         end
         if !iszerotangent(ΔR)
-            ΔR22 = view(ΔR, (p + 1):minmn, (p + 1):n)
+            ΔR22 = view(ΔR, (p + 1):minmn, (p + 1):size(R, 2))
             Δgauge = max(Δgauge, norm(ΔR22, Inf))
         end
         Δgauge ≤ gauge_atol ||
@@ -19,7 +19,7 @@ function check_qr_cotangents(Q, R, ΔQ, ΔR, minmn::Int, p::Int; gauge_atol::Rea
     return
 end
 
-function check_qr_full_cotangents(Q1, ΔQ2, ΔR, Q1dΔQ2, ; gauge_atol::Real = default_pullback_gauge_atol(ΔQ2))
+function check_qr_full_cotangents(Q1, ΔQ2, Q1dΔQ2; gauge_atol::Real = default_pullback_gauge_atol(ΔQ2))
     # in the case where A is full rank, but there are more columns in Q than in A
     # (the case of `qr_full`), there is gauge-invariant information in the
     # projection of ΔQ2 onto the column space of Q1, by virtue of Q being a unitary
@@ -113,6 +113,17 @@ function qr_pullback!(
     ΔA1 .+= ΔQ̃
     return ΔA
 end
+function qr_pullback!(
+        ΔA::Diagonal, A, QR, ΔQR;
+        rank_atol::Real = default_pullback_rank_atol(QR[2]),
+        gauge_atol::Real = default_pullback_gauge_atol(ΔQR[1])
+    )
+    ΔA_full = zero!(similar(ΔA, size(ΔA)))
+    ΔA_full = qr_pullback!(ΔA_full, A, QR, ΔQR; rank_atol, gauge_atol)
+    @assert isdiag(ΔA_full)
+    diagview(ΔA) .+= diagview(ΔA_full)
+    return ΔA
+end
 
 function check_qr_null_cotangents(N, ΔN; gauge_atol::Real = default_pullback_gauge_atol(ΔN))
     aNᴴΔN = project_antihermitian!(N' * ΔN)

test/mooncake.jl

@@ -16,7 +16,7 @@ for T in (BLASFloats..., GenericFloats...), n in (17, m, 23)
     TestSuite.seed_rng!(123)
     if CUDA.functional()
         TestSuite.test_mooncake(CuMatrix{T}, (m, n); atol = m * n * TestSuite.precision(T), rtol = m * n * TestSuite.precision(T))
-        #n == m && TestSuite.test_mooncake(Diagonal{T, CuVector{T}}, m; atol = m * TestSuite.precision(T), rtol = m * TestSuite.precision(T))
+        n == m && TestSuite.test_mooncake(Diagonal{T, CuVector{T}}, m; atol = m * TestSuite.precision(T), rtol = m * TestSuite.precision(T))
     end
     #=if AMDGPU.functional()
         TestSuite.test_mooncake(ROCMatrix{T}, (m, n); atol = m * n * TestSuite.precision(T), rtol = m * n * TestSuite.precision(T))

test/testsuite/ad_utils.jl

@@ -50,7 +50,13 @@ function make_eig_matrix(T, sz)
     D, V = eig_full(A)
     stabilize_eigvals!(diagview(D))
     Ac = V * D * inv(V)
-    return (eltype(T) <: Real) ? real(Ac) : Ac
+    Af = (eltype(T) <: Real) ? real(Ac) : Ac
+    if T <: Diagonal
+        copyto!(diagview(A), diagview(Af))
+    else
+        copyto!(A, Af)
+    end
+    return A
 end
 function make_eigh_matrix(T, sz)
     A = project_hermitian!(instantiate_matrix(T, sz))
@@ -116,7 +122,7 @@ function ad_qr_rank_deficient_compact_setup(A)
     Q1 = view(Q, 1:m, 1:r)
     Q2 = view(Q, 1:m, (r + 1):minmn)
     ΔQ2 = view(ΔQ, 1:m, (r + 1):minmn)
-    zero!(ΔQ2)
+    MatrixAlgebraKit.zero!(ΔQ2)
     ΔR = randn!(similar(A, T, minmn, n))
     view(ΔR, (r + 1):minmn, :) .= 0
     return (Q, R), (ΔQ, ΔR)
@@ -128,13 +134,13 @@ function ad_qr_rank_deficient_compact_setup(A::Diagonal)
     T = eltype(A)
     r = minmn - 5
     Ard_ = randn!(similar(A, T, m))
-    zero!(view(Ard_, (r + 1):m))
+    MatrixAlgebraKit.zero!(view(Ard_, (r + 1):m))
     Ard = Diagonal(Ard_)
     Q, R = qr_compact(Ard)
     ΔQ = Diagonal(randn!(similar(diagview(A), T, m)))
     ΔR = Diagonal(randn!(similar(diagview(A), T, m)))
-    zero!(view(diagview(ΔQ), (r + 1):m))
-    zero!(view(diagview(ΔR), (r + 1):m))
+    MatrixAlgebraKit.zero!(view(diagview(ΔQ), (r + 1):m))
+    MatrixAlgebraKit.zero!(view(diagview(ΔR), (r + 1):m))
     return (Q, R), (ΔQ, ΔR)
 end
 
@@ -206,10 +212,10 @@ end
 
 function ad_eig_full_setup(A::Diagonal)
     m, n = size(A)
-    T = eltype(A)
+    T = complex(eltype(A))
     DV = eig_full(A)
     D, V = DV
-    ΔV = Diagonal(randn!(similar(A.diag, T, m)))
+    ΔV = randn!(similar(A.diag, T, m, m))
     ΔV = remove_eiggauge_dependence!(ΔV, D, V)
     ΔD = Diagonal(randn!(similar(A.diag, T, m)))
     ΔD2 = Diagonal(randn!(similar(A.diag, T, m)))
@@ -218,17 +224,17 @@ end
 
 function ad_eig_vals_setup(A)
     m, n = size(A)
-    T = eltype(A)
+    T = complex(eltype(A))
     D = eig_vals(A)
     ΔD = randn!(similar(A, complex(T), m))
     return D, ΔD
 end
 
 function ad_eig_vals_setup(A::Diagonal)
     m, n = size(A)
-    T = eltype(A)
+    T = complex(eltype(A))
     D = eig_vals(A)
-    ΔD = randn!(similar(A.diag))
+    ΔD = randn!(similar(A.diag, T, m))
     return D, ΔD
 end