Skip to content

Commit 87124c6

Browse files
kshyattkshyatt
committed
Eigh pbs and tests
1 parent bbab562 commit 87124c6

3 files changed

Lines changed: 236 additions & 4 deletions

File tree

‎ext/PEPSKitMooncakeExt.jl‎

Lines changed: 52 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,14 @@
11
module PEPSKitMooncakeExt
22

33
using PEPSKit, TensorKit, Mooncake, MatrixAlgebraKit
4-
using PEPSKit: SVDAdjoint
4+
using PEPSKit: SVDAdjoint, EighAdjoint
55
using Mooncake: DefaultCtx, CoDual, Dual, NoRData, primal, rrule!!, arrayify, @is_primitive
66

77
_warn_pullback_truncerror(dϵ::Real; tol = MatrixAlgebraKit.defaulttol(dϵ)) =
88
abs(dϵ) ≤ tol || @warn "Pullback ignores non-zero tangents for truncation error"
99

1010
Mooncake.tangent_type(::Type{<:PEPSKit.SVDAdjoint}) = Mooncake.NoTangent
11+
Mooncake.tangent_type(::Type{<:PEPSKit.EighAdjoint}) = Mooncake.NoTangent
1112

1213
@is_primitive Mooncake.DefaultCtx Mooncake.ReverseMode Tuple{typeof(svd_trunc), TensorKit.AbstractTensorMap, SVDAdjoint}
1314
function Mooncake.rrule!!(::CoDual{typeof(MatrixAlgebraKit.svd_trunc)}, t_dt::CoDual{<:TensorKit.AbstractTensorMap}, alg_dalg::CoDual{SVDAdjoint{F, R}}) where {F, R <: PEPSKit.FullPullback}
@@ -58,4 +59,54 @@ function Mooncake.rrule!!(::CoDual{typeof(MatrixAlgebraKit.svd_trunc)}, t_dt::Co
5859
return output_codual, svd_trunc!_trunc_pullback
5960
end
6061

62+
@is_primitive Mooncake.DefaultCtx Mooncake.ReverseMode Tuple{typeof(eigh_trunc), TensorKit.AbstractTensorMap, EighAdjoint}
63+
function Mooncake.rrule!!(::CoDual{typeof(MatrixAlgebraKit.eigh_trunc)}, t_dt::CoDual{<:TensorKit.AbstractTensorMap}, alg_dalg::CoDual{EighAdjoint{F, R}}) where {F, R <: PEPSKit.FullPullback}
64+
t, dt = arrayify(t_dt)
65+
alg = primal(alg_dalg)
66+
67+
D, V = eigh_full!(t; alg.fwd_alg.alg)
68+
(D̃, Ṽ), inds = MatrixAlgebraKit.truncate(eigh_trunc!, (D, V), alg.fwd_alg.trunc)
69+
ϵ = MatrixAlgebraKit.truncation_error(diagview(D), inds)
70+
71+
DVtrunc = (D̃, Ṽ)
72+
# pack output
73+
DVtrunc_dDVtrunc = Mooncake.zero_fcodual((DVtrunc..., ϵ))
74+
75+
# define pullback
76+
dDVtrunc = last.(arrayify.(DVtrunc, Base.front(Mooncake.tangent(DVtrunc_dDVtrunc))))
77+
78+
gtol = PEPSKit._get_pullback_gauge_tol(alg.rrule_alg.verbosity)
79+
function eigh_trunc!_full_pullback((_, _, dϵ)::Tuple{NoRData, NoRData, Real})
80+
_warn_pullback_truncerror(dϵ)
81+
MatrixAlgebraKit.eigh_pullback!(dt, t, (D, V), dDVtrunc, inds; gauge_atol = gtol(dDVtrunc), degeneracy_atol = alg.rrule_alg.degeneracy_atol)
82+
MatrixAlgebraKit.zero!.(dDVtrunc) # since this is allocated in this function this is probably not required
83+
return ntuple(Returns(NoRData()), 3)
84+
end
85+
return DVtrunc_dDVtrunc, eigh_trunc!_full_pullback
86+
end
87+
88+
function Mooncake.rrule!!(::CoDual{typeof(MatrixAlgebraKit.eigh_trunc)}, t_dt::CoDual{<:TensorKit.AbstractTensorMap}, alg_dalg::CoDual{EighAdjoint{F, R}}) where {F, R <: PEPSKit.TruncPullback}
89+
t, dt = arrayify(t_dt)
90+
alg = primal(alg_dalg)
91+
92+
D, V, truncerror = eigh_trunc(t, alg)
93+
gtol = PEPSKit._get_pullback_gauge_tol(alg.rrule_alg.verbosity)
94+
output = (D, V, truncerror)
95+
output_codual = CoDual(output, Mooncake.fdata(Mooncake.zero_tangent(output)))
96+
97+
gtol = PEPSKit._get_pullback_gauge_tol(alg.rrule_alg.verbosity)
98+
function eigh_trunc!_trunc_pullback((_, _, dϵ)::Tuple{NoRData, NoRData, Real})
99+
_warn_pullback_truncerror(dϵ)
100+
Dtrunc, Vtrunc, ϵ = Mooncake.primal(output_codual)
101+
dDtrunc_, dVtrunc_, dϵ = Mooncake.tangent(output_codual)
102+
D, dD = arrayify(Dtrunc, dDtrunc_)
103+
V, dV = arrayify(Vtrunc, dVtrunc_)
104+
MatrixAlgebraKit.eigh_trunc_pullback!(dt, t, (D, V), (dD, dV); gauge_atol = gtol((dD, dV)), degeneracy_atol = alg.rrule_alg.degeneracy_atol)
105+
MatrixAlgebraKit.zero!(dD) # since this is allocated in this function this is probably not required
106+
MatrixAlgebraKit.zero!(dV) # since this is allocated in this function this is probably not required
107+
return ntuple(Returns(NoRData()), 3)
108+
end
109+
return output_codual, eigh_trunc!_trunc_pullback
110+
end
111+
61112
end

‎test/mooncake/eigh_wrapper.jl‎

Lines changed: 181 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,181 @@
1+
2+
using Test
3+
using Random
4+
using LinearAlgebra
5+
using TensorKit
6+
using Mooncake
7+
using Accessors
8+
using PEPSKit
9+
10+
using MatrixAlgebraKit: TruncatedAlgorithm, diagview
11+
12+
# Gauge-invariant loss function
13+
function lossfun(A, alg, R = randn(space(A)), trunc = notrunc())
14+
alg = @set alg.fwd_alg = TruncatedAlgorithm(alg.fwd_alg, trunc)
15+
D, V, = eigh_trunc(project_hermitian(A), alg)
16+
return real(dot(R, V * V')) + dot(D, D) # Overlap with random tensor R is gauge-invariant and differentiable
17+
end
18+
19+
dtype = ComplexF64
20+
n = 20
21+
χ = 10
22+
trunc = truncspace(ℂ^χ)
23+
rtol = 1.0e-9
24+
Random.seed!(123456789)
25+
r = randn(dtype, ℂ^n, ℂ^n)
26+
r = 0.5 * (r + r') # make r Hermitian
27+
R = randn(space(r))
28+
R = 0.5 * (R + R')
29+
30+
full_alg = EighAdjoint(; fwd_alg = (; alg = :QRIteration), rrule_alg = (; alg = :FullPullback))
31+
trunc_alg = EighAdjoint(; fwd_alg = (; alg = :QRIteration), rrule_alg = (; alg = :TruncPullback))
32+
iter_alg = EighAdjoint(; fwd_alg = (; alg = :Lanczos), rrule_alg = (; alg = :TruncPullback))
33+
34+
@testset "Non-truncated eigh" begin
35+
full_lossfun = A -> lossfun(A, full_alg, R)
36+
trunc_lossfun = A -> lossfun(A, trunc_alg, R)
37+
iter_lossfun = A -> lossfun(A, iter_alg, R)
38+
39+
full_rrule = Mooncake.build_rrule(full_lossfun, r)
40+
trunc_rrule = Mooncake.build_rrule(trunc_lossfun, r)
41+
iter_rrule = Mooncake.build_rrule(iter_lossfun, r)
42+
43+
l_full, g_full = Mooncake.value_and_gradient!!(full_rrule, full_lossfun, r)
44+
l_trunc, g_trunc = Mooncake.value_and_gradient!!(trunc_rrule, trunc_lossfun, r)
45+
l_iter, g_iter = Mooncake.value_and_gradient!!(iter_rrule, iter_lossfun, r)
46+
47+
@test l_full ≈ l_trunc ≈ l_iter
48+
@test g_full[2] ≈ g_trunc[2] rtol = rtol
49+
@test g_full[2] ≈ g_iter[2] rtol = rtol
50+
@test g_trunc[2] ≈ g_iter[2] rtol = rtol
51+
end
52+
53+
@testset "Truncated eigh with χ=$χ" begin
54+
full_lossfun = A -> lossfun(A, full_alg, R, trunc)
55+
trunc_lossfun = A -> lossfun(A, trunc_alg, R, trunc)
56+
iter_lossfun = A -> lossfun(A, iter_alg, R, trunc)
57+
58+
full_rrule = Mooncake.build_rrule(full_lossfun, r)
59+
trunc_rrule = Mooncake.build_rrule(trunc_lossfun, r)
60+
iter_rrule = Mooncake.build_rrule(iter_lossfun, r)
61+
62+
l_full, g_full = Mooncake.value_and_gradient!!(full_rrule, full_lossfun, r)
63+
l_trunc, g_trunc = Mooncake.value_and_gradient!!(trunc_rrule, trunc_lossfun, r)
64+
l_iter, g_iter = Mooncake.value_and_gradient!!(iter_rrule, iter_lossfun, r)
65+
66+
@test l_full ≈ l_trunc ≈ l_iter
67+
@test g_full[2] ≈ g_trunc[2] rtol = rtol
68+
@test g_full[2] ≈ g_iter[2] rtol = rtol
69+
@test g_trunc[2] ≈ g_iter[2] rtol = rtol
70+
end
71+
72+
@testset "Truncated eigh broadening for $(alg.rrule_alg)" for alg in [full_alg, trunc_alg]
73+
d, v = eigh_full(r)
74+
d.data[1:2:n] .= d.data[2:2:n] # make every eigenvalue two-fold degenerate
75+
r_degen = v * d * v'
76+
77+
no_broadening_no_cutoff_alg = @set alg.rrule_alg.degeneracy_atol = 1.0e-30
78+
small_broadening_alg = @set alg.rrule_alg.degeneracy_atol = 1.0e-13
79+
80+
only_lossfun = A -> lossfun(A, alg, R, trunc)
81+
no_broadening_lossfun = A -> lossfun(A, no_broadening_no_cutoff_alg, R, trunc)
82+
small_broadening_lossfun = A -> lossfun(A, small_broadening_alg, R, trunc)
83+
84+
only_rrule = Mooncake.build_rrule(only_lossfun, r_degen)
85+
no_broadening_rrule = Mooncake.build_rrule(no_broadening_lossfun, r_degen)
86+
small_broadening_rrule = Mooncake.build_rrule(small_broadening_lossfun, r_degen)
87+
88+
l_only_cutoff, g_only_cutoff = Mooncake.value_and_gradient!!(only_rrule, only_lossfun, r_degen) # cutoff sets degenerate difference to zero
89+
l_no_broadening_no_cutoff, g_no_broadening_no_cutoff = Mooncake.value_and_gradient!!( # degenerate singular value differences lead to divergent contributions
90+
no_broadening_rrule, no_broadening_lossfun, r_degen,
91+
)
92+
l_small_broadening, g_small_broadening = Mooncake.value_and_gradient!!( # broadening smoothens divergent contributions
93+
small_broadening_rrule, small_broadening_lossfun, r_degen,
94+
)
95+
96+
@test l_only_cutoff ≈ l_no_broadening_no_cutoff ≈ l_small_broadening
97+
@test norm(g_no_broadening_no_cutoff[2] - g_small_broadening[2]) > 1.0e-2 # divergences mess up the gradient
98+
@test g_only_cutoff[2] ≈ g_small_broadening[2] rtol = rtol # cutoff and broadening have similar effect
99+
end
100+
101+
symm_m, symm_n = 18, 24
102+
symm_space = Z2Space(0 => symm_m, 1 => symm_n)
103+
symm_trspace = truncspace(Z2Space(0 => symm_m ÷ 2, 1 => symm_n ÷ 3))
104+
symm_r = randn(dtype, symm_space, symm_space)
105+
symm_r = 0.5 * (symm_r + symm_r')
106+
symm_R = randn(dtype, space(symm_r))
107+
symm_R = 0.5 * (symm_R + symm_R')
108+
109+
@testset "IterEig of symmetric tensors" begin
110+
full_lossfun = A -> lossfun(A, full_alg, symm_R)
111+
trunc_lossfun = A -> lossfun(A, trunc_alg, symm_R)
112+
iter_lossfun = A -> lossfun(A, iter_alg, symm_R)
113+
114+
full_rrule = Mooncake.build_rrule(full_lossfun, symm_r)
115+
trunc_rrule = Mooncake.build_rrule(trunc_lossfun, symm_r)
116+
iter_rrule = Mooncake.build_rrule(iter_lossfun, symm_r)
117+
118+
l_full, g_full = Mooncake.value_and_gradient!!(full_rrule, full_lossfun, symm_r)
119+
l_trunc, g_trunc = Mooncake.value_and_gradient!!(trunc_rrule, trunc_lossfun, symm_r)
120+
l_iter, g_iter = Mooncake.value_and_gradient!!(iter_rrule, iter_lossfun, symm_r)
121+
122+
@test l_full ≈ l_trunc ≈ l_iter
123+
@test g_full[2] ≈ g_trunc[2] rtol = rtol
124+
@test g_full[2] ≈ g_iter[2] rtol = rtol
125+
@test g_trunc[2] ≈ g_iter[2] rtol = rtol
126+
127+
full_lossfun = A -> lossfun(A, full_alg, symm_R, symm_trspace)
128+
trunc_lossfun = A -> lossfun(A, trunc_alg, symm_R, symm_trspace)
129+
iter_lossfun = A -> lossfun(A, iter_alg, symm_R, symm_trspace)
130+
131+
full_rrule = Mooncake.build_rrule(full_lossfun, symm_r)
132+
trunc_rrule = Mooncake.build_rrule(trunc_lossfun, symm_r)
133+
iter_rrule = Mooncake.build_rrule(iter_lossfun, symm_r)
134+
135+
l_full_tr, g_full_tr = Mooncake.value_and_gradient!!(full_rrule, full_lossfun, symm_r)
136+
l_trunc_tr, g_trunc_tr = Mooncake.value_and_gradient!!(trunc_rrule, trunc_lossfun, symm_r)
137+
l_iter_tr, g_iter_tr = Mooncake.value_and_gradient!!(iter_rrule, iter_lossfun, symm_r)
138+
@test l_full_tr ≈ l_trunc_tr ≈ l_iter_tr
139+
@test g_full_tr[2] ≈ g_trunc_tr[2] rtol = rtol
140+
@test g_full_tr[2] ≈ g_iter_tr[2] rtol = rtol
141+
@test g_trunc_tr[2] ≈ g_iter_tr[2] rtol = rtol
142+
143+
iter_alg_fallback = @set iter_alg.fwd_alg.fallback_threshold = 0.4 # Do dense decomposition in one block, sparse one in the other
144+
fb_lossfun = A -> lossfun(A, iter_alg_fallback, symm_R, symm_trspace)
145+
fb_rrule = Mooncake.build_rrule(fb_lossfun, symm_r)
146+
l_iter_fb, g_iter_fb = Mooncake.value_and_gradient!!(fb_rrule, fb_lossfun, symm_r)
147+
@test l_iter_fb ≈ l_trunc_tr ≈ l_full_tr
148+
@test g_full_tr[2] ≈ g_iter_fb[2] rtol = rtol
149+
@test g_trunc_tr[2] ≈ g_iter_fb[2] rtol = rtol
150+
end
151+
#=
152+
@testset "Truncated symmetric eigh broadening for $(alg.rrule_alg)" for alg in [full_alg, trunc_alg]
153+
d, v = eigh_full(symm_r)
154+
# make every singular value in the 0-sector three-fold degenerate
155+
b0 = diagview(block(d, Z2Irrep(0)))
156+
b0[1:3:symm_m] .= b0[3:3:symm_m]
157+
b0[2:3:symm_m] .= b0[3:3:symm_m]
158+
# make every singular value in the 1-sector two-fold degenerate
159+
b1 = diagview(block(d, Z2Irrep(1)))
160+
b1[1:2:symm_n] .= b1[2:2:symm_n]
161+
symm_r_degen = v * d * v'
162+
163+
no_broadening_no_cutoff_alg = @set alg.rrule_alg.degeneracy_atol = 1.0e-30
164+
small_broadening_alg = @set alg.rrule_alg.degeneracy_atol = 1.0e-13
165+
166+
l_only_cutoff, g_only_cutoff = withgradient(
167+
A -> lossfun(A, alg, symm_R, symm_trspace), symm_r_degen
168+
) # cutoff sets degenerate difference to zero
169+
l_no_broadening_no_cutoff, g_no_broadening_no_cutoff = withgradient( # degenerate singular value differences lead to divergent contributions
170+
A -> lossfun(A, no_broadening_no_cutoff_alg, symm_R, symm_trspace),
171+
symm_r_degen,
172+
)
173+
l_small_broadening, g_small_broadening = withgradient( # broadening smoothens divergent contributions
174+
A -> lossfun(A, small_broadening_alg, symm_R, symm_trspace),
175+
symm_r_degen,
176+
)
177+
178+
@test l_only_cutoff ≈ l_no_broadening_no_cutoff ≈ l_small_broadening
179+
@test norm(g_no_broadening_no_cutoff[1] - g_small_broadening[1]) > 1.0e-2 # divergences mess up the gradient
180+
@test g_only_cutoff[1] ≈ g_small_broadening[1] rtol = rtol # cutoff and broadening have similar effect
181+
end=#

‎test/mooncake/svd_wrapper.jl‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -75,15 +75,15 @@ end
7575
no_broadening_no_cutoff_alg = @set full_alg.rrule_alg.degeneracy_atol = 1.0e-30
7676
small_broadening_alg = @set full_alg.rrule_alg.degeneracy_atol = 1.0e-13
7777

78-
full_lossfun = A -> lossfun(A, full_alg, R, trunc)
78+
only_lossfun = A -> lossfun(A, alg, R, trunc)
7979
no_broadening_lossfun = A -> lossfun(A, no_broadening_no_cutoff_alg, R, trunc)
8080
small_broadening_lossfun = A -> lossfun(A, small_broadening_alg, R, trunc)
8181

82-
full_rrule = Mooncake.build_rrule(full_lossfun, r_degen)
82+
only_rrule = Mooncake.build_rrule(only_lossfun, r_degen)
8383
no_broadening_rrule = Mooncake.build_rrule(no_broadening_lossfun, r_degen)
8484
small_broadening_rrule = Mooncake.build_rrule(small_broadening_lossfun, r_degen)
8585

86-
l_only_cutoff, g_only_cutoff = Mooncake.value_and_gradient!!(full_rrule, full_lossfun, r_degen) # cutoff sets degenerate difference to zero
86+
l_only_cutoff, g_only_cutoff = Mooncake.value_and_gradient!!(only_rrule, only_lossfun, r_degen) # cutoff sets degenerate difference to zero
8787
l_no_broadening_no_cutoff, g_no_broadening_no_cutoff = Mooncake.value_and_gradient!!( # degenerate singular value differences lead to divergent contributions
8888
no_broadening_rrule, no_broadening_lossfun, r_degen,
8989
)

0 commit comments

Comments
 (0)