diff --git a/Project.toml b/Project.toml index 4344242d5..04eb90c2c 100644 --- a/Project.toml +++ b/Project.toml @@ -48,7 +48,7 @@ Adapt = "4" CUDA = "6" ChainRulesCore = "1" Dictionaries = "0.4" -Enzyme = "0.13.157" +Enzyme = "0.13.183" EnzymeTestUtils = "0.2.8" FiniteDifferences = "0.12" GPUArrays = "11.4.1" diff --git a/ext/TensorKitEnzymeExt/TensorKitEnzymeExt.jl b/ext/TensorKitEnzymeExt/TensorKitEnzymeExt.jl index cadae496a..2095b3c52 100644 --- a/ext/TensorKitEnzymeExt/TensorKitEnzymeExt.jl +++ b/ext/TensorKitEnzymeExt/TensorKitEnzymeExt.jl @@ -15,5 +15,6 @@ include("utility.jl") include("linalg.jl") include("indexmanipulations.jl") include("tensoroperations.jl") +include("factorizations.jl") end diff --git a/ext/TensorKitEnzymeExt/factorizations.jl b/ext/TensorKitEnzymeExt/factorizations.jl new file mode 100644 index 000000000..4e6b6e962 --- /dev/null +++ b/ext/TensorKitEnzymeExt/factorizations.jl @@ -0,0 +1,69 @@ +# need these due to Enzyme choking on blocks + +for f in (:project_hermitian, :project_antihermitian) + f! = Symbol(f, :!) + @eval begin + function EnzymeRules.augmented_primal( + config::EnzymeRules.RevConfigWidth{1}, + func::Const{typeof($f!)}, + ::Type{RT}, + A::Annotation{<:AbstractTensorMap}, + arg::Annotation{<:AbstractTensorMap}, + alg::Const, + ) where {RT} + $f!(A.val, arg.val, alg.val) + primal = EnzymeRules.needs_primal(config) ? arg.val : nothing + shadow = EnzymeRules.needs_shadow(config) ? arg.dval : nothing + cache = nothing + return EnzymeRules.AugmentedReturn(primal, shadow, cache) + end + function EnzymeRules.reverse( + config::EnzymeRules.RevConfigWidth{1}, + func::Const{typeof($f!)}, + ::Type{RT}, + cache, + A::Annotation{<:AbstractTensorMap}, + arg::Annotation{<:AbstractTensorMap}, + alg::Const, + ) where {RT} + if !isa(A, Const) && !isa(arg, Const) + $f!(arg.dval, arg.dval, alg.val) + if A.dval !== arg.dval + A.dval .+= arg.dval + make_zero!(arg.dval) + end + end + return (nothing, nothing, nothing) + end + function EnzymeRules.augmented_primal( + config::EnzymeRules.RevConfigWidth{1}, + func::Const{typeof($f)}, + ::Type{RT}, + A::Annotation{<:AbstractTensorMap}, + alg::Const, + ) where {RT} + ret = $f(A.val, alg.val) + dret = make_zero(ret) + primal = EnzymeRules.needs_primal(config) ? ret : nothing + shadow = EnzymeRules.needs_shadow(config) ? dret : nothing + cache = dret + return EnzymeRules.AugmentedReturn(primal, shadow, cache) + end + function EnzymeRules.reverse( + config::EnzymeRules.RevConfigWidth{1}, + func::Const{typeof($f)}, + ::Type{RT}, + cache, + A::Annotation{<:AbstractTensorMap}, + alg::Const, + ) where {RT} + dret = cache + if !isa(A, Const) + $f!(dret, dret, alg.val) + add!(A.dval, dret) + end + make_zero!(dret) + return (nothing, nothing) + end + end +end diff --git a/ext/TensorKitEnzymeExt/utility.jl b/ext/TensorKitEnzymeExt/utility.jl index dd99d96e8..d03848c81 100644 --- a/ext/TensorKitEnzymeExt/utility.jl +++ b/ext/TensorKitEnzymeExt/utility.jl @@ -135,6 +135,7 @@ end @inline EnzymeRules.inactive(::typeof(TensorKit.insertleftunit), ::HomSpace, ::Any) = nothing @inline EnzymeRules.inactive(::typeof(TensorKit.insertrightunit), ::HomSpace, ::Any) = nothing @inline EnzymeRules.inactive(::typeof(TensorKit.removeunit), ::HomSpace, ::Any) = nothing +@inline EnzymeRules.inactive(::typeof(TensorKit.infimum), ::Any, ::Any) = nothing @inline EnzymeRules.inactive(::typeof(TensorKit.sectorstructure), ::Any) = nothing @inline EnzymeRules.inactive(::typeof(TensorKit.degeneracystructure), ::Any) = nothing @inline EnzymeRules.inactive(::typeof(TensorKit.select), s::HomSpace, i::Index2Tuple) = nothing diff --git a/src/factorizations/diagonal.jl b/src/factorizations/diagonal.jl index dae550ea1..49261f24e 100644 --- a/src/factorizations/diagonal.jl +++ b/src/factorizations/diagonal.jl @@ -3,6 +3,7 @@ _repack_diagonal(d::DiagonalTensorMap) = Diagonal(d.data) _repack_diagonal(d::SectorVector) = Diagonal(parent(d)) +MAK.diagonal(t::SectorVector) = DiagonalTensorMap(t) MAK.diagview(t::DiagonalTensorMap) = SectorVector(t.data, TensorKit.diagonalblockstructure(space(t))) for f in ( diff --git a/src/factorizations/pullbacks.jl b/src/factorizations/pullbacks.jl index b910a184c..72ad94008 100644 --- a/src/factorizations/pullbacks.jl +++ b/src/factorizations/pullbacks.jl @@ -11,6 +11,16 @@ for pullback! in ( end return Δt end + @eval function MAK.$pullback!( + Δt::AbstractTensorMap, ::Nothing, F, ΔF; kwargs... + ) + foreachblock(Δt) do c, (Δb,) + Fc = block.(F, Ref(c)) + ΔFc = block.(ΔF, Ref(c)) + return MAK.$pullback!(Δb, nothing, Fc, ΔFc; kwargs...) + end + return Δt + end end for pullback! in (:qr_null_pullback!, :lq_null_pullback!) @eval function MAK.$pullback!( @@ -41,6 +51,33 @@ for pullback! in (:svd_pullback!, :eig_pullback!, :eigh_pullback!) end return Δt end + @eval function MAK.$pullback!( + Δt::AbstractTensorMap, ::Nothing, F, ΔF, inds; kwargs... + ) + foreachblock(Δt) do c, (Δb,) + haskey(inds, c) || return nothing + ind = inds[c] + Fc = block.(F, Ref(c)) + ΔFc = map(ΔFc -> isnothing(ΔFc) ? nothing : block(ΔFc, c), ΔF) + return MAK.$pullback!(Δb, nothing, Fc, ΔFc, ind; kwargs...) + end + return Δt + end + @eval function MAK.$pullback!( + Δt::AbstractTensorMap, t::AbstractTensorMap, F, ΔF, ::Colon; kwargs... + ) + return MAK.$pullback!(Δt, t, F, ΔF, _notrunc_ind(t); kwargs...) + end + @eval function MAK.$pullback!( + Δt::AbstractTensorMap, ::Nothing, F, ΔF; kwargs... + ) + return MAK.$pullback!(Δt, nothing, F, ΔF, _notrunc_ind(Δt); kwargs...) + end + @eval function MAK.$pullback!( + Δt::AbstractTensorMap, ::Nothing, F, ΔF, ::Colon; kwargs... + ) + return MAK.$pullback!(Δt, nothing, F, ΔF, _notrunc_ind(Δt); kwargs...) + end end for pullback_trunc! in (:svd_trunc_pullback!, :eig_trunc_pullback!, :eigh_trunc_pullback!) @@ -97,3 +134,7 @@ function MAK.remove_svd_gauge_dependence!( end return ΔU, ΔVᴴ end + +MAK.has_equal_storage(A::AbstractTensorMap, B::AbstractTensorMap) = A === B +MAK.has_equal_storage(A::AbstractTensorMap, B::SectorVector) = false +MAK.has_equal_storage(A::SectorVector, B::AbstractTensorMap) = false diff --git a/test/enzyme-factorizations/factorizations.jl b/test/enzyme-factorizations/factorizations.jl new file mode 100644 index 000000000..d2f4e66ef --- /dev/null +++ b/test/enzyme-factorizations/factorizations.jl @@ -0,0 +1,123 @@ +using Test, TestExtras +using TensorKit +using TensorOperations +using MatrixAlgebraKit +using MatrixAlgebraKit: remove_svd_gauge_dependence! +using MatrixAlgebraKit: remove_eig_gauge_dependence! +using MatrixAlgebraKit: remove_eigh_gauge_dependence! +using MatrixAlgebraKit: remove_lq_gauge_dependence!, remove_lq_null_gauge_dependence! +using MatrixAlgebraKit: remove_qr_gauge_dependence!, remove_qr_null_gauge_dependence! +using Enzyme, EnzymeTestUtils +using Random + +spacelist = ad_spacelist(fast_tests) +eltypes = (Float64, ComplexF64) + +@timedtestset "Enzyme - Factorizations: $(TensorKit.type_repr(sectortype(eltype(V)))) ($T)" for V in spacelist, T in eltypes, t in (randn(T, V[1] ⊗ V[2] ← V[1] ⊗ V[2]), randn(T, V[1] ⊗ V[2] ← (V[3] ⊗ V[4] ⊗ V[5])')) + atol = default_tol(T) + rtol = default_tol(T) + + @testset "SVD" begin + S = svd_vals(t) + EnzymeTestUtils.test_reverse(svd_vals, Duplicated, (t, Duplicated); atol, rtol) + + USVᴴ = svd_full(t) + ΔUSVᴴ = EnzymeTestUtils.rand_tangent(USVᴴ) + remove_svd_gauge_dependence!(ΔUSVᴴ[1], ΔUSVᴴ[3], USVᴴ...) + EnzymeTestUtils.test_reverse(svd_full, Duplicated, (t, Duplicated); output_tangent = ΔUSVᴴ, atol, rtol) + + USVᴴ = svd_compact(t) + ΔUSVᴴ = EnzymeTestUtils.rand_tangent(USVᴴ) + remove_svd_gauge_dependence!(ΔUSVᴴ[1], ΔUSVᴴ[3], USVᴴ...) + EnzymeTestUtils.test_reverse(svd_compact, Duplicated, (t, Duplicated); output_tangent = ΔUSVᴴ, atol, rtol) + + V_trunc = spacetype(t)(c => min(size(b)...) ÷ 2 for (c, b) in blocks(t)) + trunc = truncspace(V_trunc) + alg = MatrixAlgebraKit.select_algorithm(svd_trunc_no_error, t, nothing; trunc) + USVᴴtrunc = svd_trunc_no_error(t, alg) + ΔUSVᴴtrunc = EnzymeTestUtils.rand_tangent(USVᴴtrunc) + remove_svd_gauge_dependence!(ΔUSVᴴtrunc[1], ΔUSVᴴtrunc[3], USVᴴtrunc...) + EnzymeTestUtils.test_reverse(svd_trunc_no_error, Duplicated, (t, Duplicated), (alg, Const); output_tangent = ΔUSVᴴtrunc, atol, rtol) + end + + @testset "LQ" begin + EnzymeTestUtils.test_reverse(lq_compact, Duplicated, (t, Duplicated); atol, rtol) + + # lq_full/lq_null requires being careful with gauges + LQ = lq_full(t) + ΔLQ = EnzymeTestUtils.rand_tangent(LQ) + remove_lq_gauge_dependence!(ΔLQ..., t, LQ...) + EnzymeTestUtils.test_reverse(lq_full, Duplicated, (t, Duplicated); output_tangent = ΔLQ, atol, rtol) + + Nᴴ = lq_null(t) + Q = lq_compact(t)[2] + ΔNᴴ = EnzymeTestUtils.rand_tangent(Nᴴ) + remove_lq_null_gauge_dependence!(ΔNᴴ, Q, Nᴴ) + EnzymeTestUtils.test_reverse(lq_null, Duplicated, (t, Duplicated); output_tangent = ΔNᴴ, atol, rtol) + end + + @testset "QR" begin + # qr_full/qr_null requires being careful with gauges + QR = qr_full(t) + ΔQR = EnzymeTestUtils.rand_tangent(QR) + remove_qr_gauge_dependence!(ΔQR..., t, QR...) + EnzymeTestUtils.test_reverse(qr_full, Duplicated, (t, Duplicated); output_tangent = ΔQR, atol, rtol) + + N = qr_null(t) + Q = qr_compact(t)[1] + ΔN = EnzymeTestUtils.rand_tangent(N) + remove_qr_null_gauge_dependence!(ΔN, t, N) + EnzymeTestUtils.test_reverse(qr_null, Duplicated, (t, Duplicated); atol, rtol, output_tangent = ΔN) + end +end + +@timedtestset "Enzyme - Factorizations (EIGH/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])) + atol = default_tol(T) + rtol = default_tol(T) + + @testset "EIG" begin + DV = eig_full(t) + ΔDV = EnzymeTestUtils.rand_tangent(DV) + remove_eig_gauge_dependence!(ΔDV[2], DV...) + EnzymeTestUtils.test_reverse(eig_full, Duplicated, (t, Duplicated); output_tangent = ΔDV, atol, rtol) + + D = eig_vals(t) + EnzymeTestUtils.test_reverse(eig_vals, Duplicated, (t, Duplicated); atol, rtol) + + V_trunc = spacetype(t)(c => min(size(b)...) ÷ 2 for (c, b) in blocks(t)) + trunc = truncspace(V_trunc) + alg = MatrixAlgebraKit.select_algorithm(eig_trunc_no_error, t, nothing; trunc) + DVtrunc = eig_trunc_no_error(t, alg) + ΔDVtrunc = EnzymeTestUtils.rand_tangent(DVtrunc) + remove_eig_gauge_dependence!(ΔDVtrunc[2], DVtrunc...) + EnzymeTestUtils.test_reverse(eig_trunc_no_error, Duplicated, (t, Duplicated), (alg, Const); output_tangent = ΔDVtrunc, atol, rtol) + end + + @testset "EIGH" begin + th = project_hermitian(t) + DV = eigh_full(th) + ΔDV = EnzymeTestUtils.rand_tangent(DV) + remove_eigh_gauge_dependence!(ΔDV[2], DV...) + proj_eigh_full(t) = eigh_full(project_hermitian(t)) + EnzymeTestUtils.test_reverse(proj_eigh_full, Duplicated, (th, Duplicated); output_tangent = ΔDV, atol, rtol) + + D = eigh_vals(th) + EnzymeTestUtils.test_reverse(eigh_vals ∘ project_hermitian, Duplicated, (th, Duplicated); atol, rtol) + + V_trunc = spacetype(th)(c => min(size(b)...) ÷ 2 for (c, b) in blocks(t)) + trunc = truncspace(V_trunc) + alg = MatrixAlgebraKit.select_algorithm(eigh_trunc_no_error, th, nothing; trunc) + DVtrunc = eigh_trunc_no_error(th, alg) + ΔDVtrunc = EnzymeTestUtils.rand_tangent(DVtrunc) + remove_eigh_gauge_dependence!(ΔDVtrunc[2], DVtrunc...) + proj_eigh(t, alg) = eigh_trunc_no_error(project_hermitian(t), alg) + EnzymeTestUtils.test_reverse(proj_eigh, Duplicated, (th, Duplicated), (alg, Const); output_tangent = ΔDVtrunc, atol, rtol) + end + + @testset "Projections" begin + EnzymeTestUtils.test_reverse(project_hermitian, Duplicated, (t, Duplicated); atol, rtol) + EnzymeTestUtils.test_reverse(project_antihermitian, Duplicated, (t, Duplicated); atol, rtol) + EnzymeTestUtils.test_reverse(project_hermitian!, Duplicated, (t, Duplicated); atol, rtol) + EnzymeTestUtils.test_reverse(project_antihermitian!, Duplicated, (t, Duplicated); atol, rtol) + end +end