Skip to content

Commit e0a9f0c

Browse files
kshyattkshyatt
committed
Add Enzyme rules for factorizations
1 parent 0988e79 commit e0a9f0c

11 files changed

Lines changed: 289 additions & 0 deletions

File tree

‎Project.toml‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,9 @@ TensorKitFiniteDifferencesExt = "FiniteDifferences"
4242
TensorKitGPUArraysExt = "GPUArrays"
4343
TensorKitMooncakeExt = "Mooncake"
4444

45+
[sources]
46+
MatrixAlgebraKit = {url = "https://github.com/quantumkithub/matrixalgebrakit.jl", rev = "main"}
47+
4548
[compat]
4649
AMDGPU = "2"
4750
Adapt = "4"

‎ext/TensorKitEnzymeExt/TensorKitEnzymeExt.jl‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,5 +13,6 @@ using Random: AbstractRNG
1313
include("utility.jl")
1414
include("linalg.jl")
1515
include("indexmanipulations.jl")
16+
include("factorizations.jl")
1617

1718
end

‎ext/TensorKitEnzymeExt/factorizations.jl‎

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
# need these due to Enzyme choking on blocks
2+
3+
for f in (:project_hermitian, :project_antihermitian)
4+
f! = Symbol(f, :!)
5+
@eval begin
6+
function EnzymeRules.augmented_primal(
7+
config::EnzymeRules.RevConfigWidth{1},
8+
func::Const{typeof($f!)},
9+
::Type{RT},
10+
A::Annotation{<:AbstractTensorMap},
11+
arg::Annotation{<:AbstractTensorMap},
12+
alg::Const,
13+
) where {RT}
14+
$f!(A.val, arg.val, alg.val)
15+
primal = EnzymeRules.needs_primal(config) ? arg.val : nothing
16+
shadow = EnzymeRules.needs_shadow(config) ? arg.dval : nothing
17+
cache = nothing
18+
return EnzymeRules.AugmentedReturn(primal, shadow, cache)
19+
end
20+
function EnzymeRules.reverse(
21+
config::EnzymeRules.RevConfigWidth{1},
22+
func::Const{typeof($f!)},
23+
::Type{RT},
24+
cache,
25+
A::Annotation{<:AbstractTensorMap},
26+
arg::Annotation{<:AbstractTensorMap},
27+
alg::Const,
28+
) where {RT}
29+
if !isa(A, Const) && !isa(arg, Const)
30+
$f!(arg.dval, arg.dval, alg.val)
31+
if A.dval !== arg.dval
32+
A.dval .+= arg.dval
33+
make_zero!(arg.dval)
34+
end
35+
end
36+
return (nothing, nothing, nothing)
37+
end
38+
function EnzymeRules.augmented_primal(
39+
config::EnzymeRules.RevConfigWidth{1},
40+
func::Const{typeof($f)},
41+
::Type{RT},
42+
A::Annotation{<:AbstractTensorMap},
43+
alg::Const,
44+
) where {RT}
45+
ret = $f(A.val, alg.val)
46+
dret = make_zero(ret)
47+
primal = EnzymeRules.needs_primal(config) ? ret : nothing
48+
shadow = EnzymeRules.needs_shadow(config) ? dret : nothing
49+
cache = dret
50+
return EnzymeRules.AugmentedReturn(primal, shadow, cache)
51+
end
52+
function EnzymeRules.reverse(
53+
config::EnzymeRules.RevConfigWidth{1},
54+
func::Const{typeof($f)},
55+
::Type{RT},
56+
cache,
57+
A::Annotation{<:AbstractTensorMap},
58+
alg::Const,
59+
) where {RT}
60+
dret = cache
61+
if !isa(A, Const)
62+
$f!(dret, dret, alg.val)
63+
add!(A.dval, dret)
64+
end
65+
make_zero!(dret)
66+
return (nothing, nothing)
67+
end
68+
end
69+
end

‎ext/TensorKitEnzymeExt/utility.jl‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ pullback_dC!(ΔC, β::Number) = scale!(ΔC, conj(β))
1919
@inline EnzymeRules.inactive_type(::Type{<:TensorKit.GenericTreeTransformer}) = true
2020
@inline EnzymeRules.inactive_type(::Type{<:TensorKit.VectorSpace}) = true
2121

22+
@inline EnzymeRules.inactive(::typeof(TensorKit.infimum), ::Any, ::Any) = nothing
2223
@inline EnzymeRules.inactive(::typeof(TensorKit.sectorstructure), ::Any) = nothing
2324
@inline EnzymeRules.inactive(::typeof(TensorKit.degeneracystructure), ::Any) = nothing
2425
@inline EnzymeRules.inactive(::typeof(TensorKit.select), s::HomSpace, i::Index2Tuple) = nothing

‎src/factorizations/pullbacks.jl‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,16 @@ for pullback! in (
1111
end
1212
return Δt
1313
end
14+
@eval function MAK.$pullback!(
15+
Δt::AbstractTensorMap, ::Nothing, F, ΔF; kwargs...
16+
)
17+
foreachblock(Δt) do c, (Δb,)
18+
Fc = block.(F, Ref(c))
19+
ΔFc = block.(ΔF, Ref(c))
20+
return MAK.$pullback!(Δb, nothing, Fc, ΔFc; kwargs...)
21+
end
22+
return Δt
23+
end
1424
end
1525
for pullback! in (:qr_null_pullback!, :lq_null_pullback!)
1626
@eval function MAK.$pullback!(
@@ -41,6 +51,28 @@ for pullback! in (:svd_pullback!, :eig_pullback!, :eigh_pullback!)
4151
end
4252
return Δt
4353
end
54+
@eval function MAK.$pullback!(
55+
Δt::AbstractTensorMap, ::Nothing, F, ΔF, inds; kwargs...
56+
)
57+
foreachblock(Δt) do c, (Δb,)
58+
haskey(inds, c) || return nothing
59+
ind = inds[c]
60+
Fc = block.(F, Ref(c))
61+
ΔFc = block.(ΔF, Ref(c))
62+
return MAK.$pullback!(Δb, nothing, Fc, ΔFc, ind; kwargs...)
63+
end
64+
return Δt
65+
end
66+
@eval function MAK.$pullback!(
67+
Δt::AbstractTensorMap, t::AbstractTensorMap, F, ΔF, ::Colon; kwargs...
68+
)
69+
return MAK.$pullback!(Δt, t, F, ΔF, _notrunc_ind(t); kwargs...)
70+
end
71+
@eval function MAK.$pullback!(
72+
Δt, ::Nothing, F, ΔF; kwargs...
73+
)
74+
return MAK.$pullback!(Δt, nothing, F, ΔF, _notrunc_ind(Δt); kwargs...)
75+
end
4476
end
4577

4678
for pullback_trunc! in (:svd_trunc_pullback!, :eig_trunc_pullback!, :eigh_trunc_pullback!)
@@ -97,3 +129,7 @@ function MAK.remove_svd_gauge_dependence!(
97129
end
98130
return ΔU, ΔVᴴ
99131
end
132+
133+
MAK.has_equal_storage(A::AbstractTensorMap, B::AbstractTensorMap) = A === B
134+
MAK.has_equal_storage(A::AbstractTensorMap, B::SectorVector) = false
135+
MAK.has_equal_storage(A::SectorVector, B::AbstractTensorMap) = false

‎test/enzyme-factorizations/eig.jl‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
using Test, TestExtras
2+
using TensorKit
3+
using TensorOperations
4+
using VectorInterface: Zero, One
5+
using MatrixAlgebraKit
6+
using MatrixAlgebraKit: remove_eig_gauge_dependence!
7+
using Enzyme, EnzymeTestUtils
8+
using Random
9+
10+
spacelist = ad_spacelist(fast_tests)
11+
eltypes = (Float64, ComplexF64)
12+
13+
@timedtestset "Enzyme - Factorizations (EIG): $(TensorKit.type_repr(sectortype(eltype(V)))) ($T)" for V in spacelist, T in eltypes, t in (randn(T, V[1] ← V[1]), rand(T, V[1] ⊗ V[2] ← V[1] ⊗ V[2]))
14+
atol = default_tol(T)
15+
rtol = default_tol(T)
16+
DV = eig_full(t)
17+
ΔDV = EnzymeTestUtils.rand_tangent(DV)
18+
remove_eig_gauge_dependence!(ΔDV[2], DV...)
19+
EnzymeTestUtils.test_reverse(eig_full, Duplicated, (t, Duplicated); output_tangent = ΔDV, atol, rtol)
20+
21+
#D = eig_vals(t)
22+
#EnzymeTestUtils.test_reverse(eig_vals, Duplicated, (t, Duplicated); atol, rtol)
23+
24+
V_trunc = spacetype(t)(c => min(size(b)...) ÷ 2 for (c, b) in blocks(t))
25+
trunc = truncspace(V_trunc)
26+
alg = MatrixAlgebraKit.select_algorithm(eig_trunc_no_error, t, nothing; trunc)
27+
DVtrunc = eig_trunc_no_error(t, alg)
28+
ΔDVtrunc = EnzymeTestUtils.rand_tangent(DVtrunc)
29+
remove_eig_gauge_dependence!(ΔDVtrunc[2], DVtrunc...)
30+
EnzymeTestUtils.test_reverse(eig_trunc_no_error, Duplicated, (t, Duplicated), (alg, Const); output_tangent = ΔDVtrunc, atol, rtol)
31+
end

‎test/enzyme-factorizations/eigh.jl‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
using Test, TestExtras
2+
using TensorKit
3+
using TensorOperations
4+
using MatrixAlgebraKit
5+
using MatrixAlgebraKit: remove_eigh_gauge_dependence!
6+
using Enzyme, EnzymeTestUtils
7+
using Random
8+
9+
spacelist = ad_spacelist(fast_tests)
10+
eltypes = (Float64, ComplexF64)
11+
12+
@timedtestset "Enzyme - Factorizations (EIGH): $(TensorKit.type_repr(sectortype(eltype(V)))) ($T)" for V in spacelist, T in eltypes, t in (randn(T, V[1] ← V[1]), rand(T, V[1] ⊗ V[2] ← V[1] ⊗ V[2]))
13+
atol = default_tol(T)
14+
rtol = default_tol(T)
15+
th = project_hermitian(t)
16+
DV = eigh_full(th)
17+
ΔDV = EnzymeTestUtils.rand_tangent(DV)
18+
remove_eigh_gauge_dependence!(ΔDV[2], DV...)
19+
proj_eigh_full(t) = eigh_full(project_hermitian(t))
20+
EnzymeTestUtils.test_reverse(proj_eigh_full, Duplicated, (th, Duplicated); output_tangent = ΔDV, atol, rtol)
21+
22+
#D = eigh_vals(th)
23+
#EnzymeTestUtils.test_reverse(eigh_vals ∘ project_hermitian, Duplicated, (th, Duplicated); atol, rtol)
24+
25+
V_trunc = spacetype(th)(c => min(size(b)...) ÷ 2 for (c, b) in blocks(t))
26+
trunc = truncspace(V_trunc)
27+
alg = MatrixAlgebraKit.select_algorithm(eigh_trunc_no_error, th, nothing; trunc)
28+
DVtrunc = eigh_trunc_no_error(th, alg)
29+
ΔDVtrunc = EnzymeTestUtils.rand_tangent(DVtrunc)
30+
remove_eigh_gauge_dependence!(ΔDVtrunc[2], DVtrunc...)
31+
proj_eigh(t, alg) = eigh_trunc_no_error(project_hermitian(t), alg)
32+
EnzymeTestUtils.test_reverse(proj_eigh, Duplicated, (th, Duplicated), (alg, Const); output_tangent = ΔDVtrunc, atol, rtol)
33+
end

‎test/enzyme-factorizations/lq.jl‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
using Test, TestExtras
2+
using TensorKit
3+
using TensorOperations
4+
using MatrixAlgebraKit
5+
using MatrixAlgebraKit: remove_lq_gauge_dependence!, remove_lq_null_gauge_dependence!
6+
using Enzyme, EnzymeTestUtils
7+
using Random
8+
9+
spacelist = ad_spacelist(fast_tests)
10+
eltypes = (Float64, ComplexF64)
11+
12+
@timedtestset "Enzyme - Factorizations (LQ): $(TensorKit.type_repr(sectortype(eltype(V)))) ($T)" for V in spacelist, T in eltypes, A in (randn(T, V[1] ⊗ V[2] ← V[1] ⊗ V[2]), randn(T, V[1] ⊗ V[2] ← (V[3] ⊗ V[4] ⊗ V[5])'))
13+
atol = default_tol(T)
14+
rtol = default_tol(T)
15+
EnzymeTestUtils.test_reverse(lq_compact, Duplicated, (A, Duplicated); atol, rtol)
16+
17+
# lq_full/lq_null requires being careful with gauges
18+
LQ = lq_full(A)
19+
ΔLQ = EnzymeTestUtils.rand_tangent(LQ)
20+
remove_lq_gauge_dependence!(ΔLQ..., A, LQ...)
21+
EnzymeTestUtils.test_reverse(lq_full, Duplicated, (A, Duplicated); output_tangent = ΔLQ, atol, rtol)
22+
23+
Nᴴ = lq_null(A)
24+
Q = lq_compact(A)[2]
25+
ΔNᴴ = EnzymeTestUtils.rand_tangent(Nᴴ)
26+
remove_lq_null_gauge_dependence!(ΔNᴴ, Q, Nᴴ)
27+
EnzymeTestUtils.test_reverse(lq_null, Duplicated, (A, Duplicated); output_tangent = ΔNᴴ, atol, rtol)
28+
end

‎test/enzyme-factorizations/projections.jl‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
using Test, TestExtras
2+
using TensorKit
3+
using TensorOperations
4+
using MatrixAlgebraKit
5+
using Enzyme, EnzymeTestUtils
6+
using Random
7+
8+
spacelist = ad_spacelist(fast_tests)
9+
eltypes = (Float64, ComplexF64)
10+
11+
@timedtestset "Enzyme - Factorizations (PROJECTIONS): $(TensorKit.type_repr(sectortype(eltype(V)))) ($T)" for V in spacelist, T in eltypes, t in (randn(T, V[1] ← V[1]), rand(T, V[1] ⊗ V[2] ← V[1] ⊗ V[2]))
12+
atol = default_tol(T)
13+
rtol = default_tol(T)
14+
EnzymeTestUtils.test_reverse(project_hermitian, Duplicated, (t, Duplicated); atol, rtol)
15+
EnzymeTestUtils.test_reverse(project_antihermitian, Duplicated, (t, Duplicated); atol, rtol)
16+
EnzymeTestUtils.test_reverse(project_hermitian!, Duplicated, (t, Duplicated); atol, rtol)
17+
EnzymeTestUtils.test_reverse(project_antihermitian!, Duplicated, (t, Duplicated); atol, rtol)
18+
end

‎test/enzyme-factorizations/qr.jl‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
using Test, TestExtras
2+
using TensorKit
3+
using TensorOperations
4+
using VectorInterface: Zero, One
5+
using MatrixAlgebraKit
6+
using MatrixAlgebraKit: remove_qr_gauge_dependence!, remove_qr_null_gauge_dependence!
7+
using Enzyme, EnzymeTestUtils
8+
using Random
9+
10+
spacelist = ad_spacelist(fast_tests)
11+
eltypes = (Float64, ComplexF64)
12+
13+
@timedtestset "Enzyme - Factorizations (QR): $(TensorKit.type_repr(sectortype(eltype(V)))) ($T)" for V in spacelist, T in eltypes, A in (randn(T, V[1] ⊗ V[2] ← V[1] ⊗ V[2]), randn(T, V[1] ⊗ V[2] ⊗ V[3] ← (V[4] ⊗ V[5])'))
14+
atol = default_tol(T)
15+
rtol = default_tol(T)
16+
17+
EnzymeTestUtils.test_reverse(qr_compact, Duplicated, (A, Duplicated); atol, rtol)
18+
19+
# qr_full/qr_null requires being careful with gauges
20+
QR = qr_full(A)
21+
ΔQR = EnzymeTestUtils.rand_tangent(QR)
22+
remove_qr_gauge_dependence!(ΔQR..., A, QR...)
23+
EnzymeTestUtils.test_reverse(qr_full, Duplicated, (A, Duplicated); output_tangent = ΔQR, atol, rtol)
24+
25+
N = qr_null(A)
26+
Q = qr_compact(A)[1]
27+
ΔN = EnzymeTestUtils.rand_tangent(N)
28+
remove_qr_null_gauge_dependence!(ΔN, A, N)
29+
EnzymeTestUtils.test_reverse(qr_null, Duplicated, (A, Duplicated); atol, rtol, output_tangent = ΔN)
30+
end

0 commit comments

Comments
 (0)