Add Enzyme rules for factorizations

.github/workflows/CI.yml

@@ -22,6 +22,6 @@ jobs:
     with:
       fast: "${{ github.event.pull_request.draft == true }}"
       exclude: '["cuda", "amd"]'
-      timeout-minutes: 120
+      timeout-minutes: 150
     secrets:
       CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}

Project.toml

@@ -51,7 +51,7 @@ Adapt = "4"
 CUDA = "6"
 ChainRulesCore = "1"
 Dictionaries = "0.4"
-Enzyme = "0.13.207"
+Enzyme = "0.13.216"
 EnzymeTestUtils = "0.2.8"
 FiniteDifferences = "0.12"
 GPUArrays = "11.4.1"

ext/TensorKitEnzymeExt/utility.jl

@@ -152,6 +152,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

ext/TensorKitEnzymeTestUtilsExt.jl

@@ -43,6 +43,33 @@ function EnzymeTestUtils.to_vec(t::TensorKit.DiagonalTensorMap, seen_vecs::Enzym
     parent_vec, parent_t = to_vec(TensorMap(t), seen_vecs)
     return parent_vec, TensorKit.DiagonalTensorMap ∘ parent_t
 end
+function EnzymeTestUtils.to_vec(v::TensorKit.SectorVector, seen_vecs::EnzymeTestUtils.AliasDict)
+    has_seen = haskey(seen_vecs, v)
+    is_const = Enzyme.Compiler.guaranteed_const(Core.Typeof(v))
+    if has_seen || is_const
+        v_vec = Float32[]
+    else
+        vec_of_vecs = [b * TensorKit.sqrtdim(c) for (c, b) in pairs(v)]
+        v_vec, back = to_vec(vec_of_vecs)
+        seen_vecs[v] = v_vec
+    end
+    function SectorVector_from_vec(v_vec_new::AbstractVector, seen_xs::EnzymeTestUtils.AliasDict)
+        if xor(has_seen, haskey(seen_xs, v))
+            throw(ErrorException("Arrays must be reconstructed in the same order as they are vectorized."))
+        end
+        has_seen && return seen_xs[v]
+        is_const && return v
+
+        v_new = similar(v)
+        vvec_of_vecs = back(v_vec_new)
+        for (i, (c, b)) in enumerate(pairs(v_new))
+            scale!(b, vvec_of_vecs[i], TensorKit.invsqrtdim(c))
+        end
+        seen_xs[v] = v_new
+        return v_new
+    end
+    return v_vec, SectorVector_from_vec
+end
 
 # generate random tangents for testing
 function EnzymeTestUtils.rand_tangent(rng, t::TensorMap)

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 (

src/factorizations/factorizations.jl

@@ -32,6 +32,7 @@ include("truncation.jl")
 include("adjoint.jl")
 include("diagonal.jl")
 include("pullbacks.jl")
+include("pushforwards.jl")
 
 TensorKit.one!(A::AbstractMatrix) = MatrixAlgebraKit.one!(A)
 

src/factorizations/pullbacks.jl

@@ -1,13 +1,16 @@
+nothing_or_block(x, c) = isnothing(x) ? x : block(x, c)
+nothing_or_block(x::Diagonal, c) = block(MAK.diagonal(parent(x)), c)
+
 for pullback! in (
         :qr_pullback!, :lq_pullback!, :left_polar_pullback!, :right_polar_pullback!,
     )
     @eval function MAK.$pullback!(
-            Δt::AbstractTensorMap, t::AbstractTensorMap, F, ΔF; kwargs...
+            Δt::AbstractTensorMap, t, F, ΔF; kwargs...
         )
-        foreachblock(Δt, t) do c, (Δb, b)
+        foreachblock(Δt) do c, (Δb,)
             Fc = block.(F, Ref(c))
             ΔFc = block.(ΔF, Ref(c))
-            return MAK.$pullback!(Δb, b, Fc, ΔFc; kwargs...)
+            return MAK.$pullback!(Δb, nothing, Fc, ΔFc; kwargs...)
         end
         return Δt
     end
@@ -26,26 +29,47 @@ for pullback! in (:qr_null_pullback!, :lq_null_pullback!)
 end
 _notrunc_ind(t) = SectorDict(c => Colon() for c in blocksectors(t))
 
+for pullback! in (:eig_vals_pullback!, :eigh_vals_pullback!)
+    @eval function MAK.$pullback!(
+            Δt::AbstractTensorMap, t, DV::Tuple{Diagonal, <:AbstractTensorMap}, ΔD, inds;
+            kwargs...
+        )
+        return MAK.$pullback!(Δt, t, (MAK.diagonal(parent(DV[1])), DV[2]), ΔD, inds; kwargs...)
+    end
+end
+function MAK.svd_vals_pullback!(
+        Δt::AbstractTensorMap, t, USVᴴ::Tuple{<:AbstractTensorMap, Diagonal, <:AbstractTensorMap}, ΔS, ind;
+        kwargs...
+    )
+    return MAK.svd_vals_pullback!(Δt, t, (USVᴴ[1], MAK.diagonal(parent(USVᴴ[2])), USVᴴ[3]), ΔS, ind; kwargs...)
+end
+
 for pullback! in (:svd_pullback!, :eig_pullback!, :eigh_pullback!)
     @eval function MAK.$pullback!(
-            Δt::AbstractTensorMap, t::AbstractTensorMap, F, ΔF, inds = _notrunc_ind(t);
+            Δt::AbstractTensorMap, t, F, ΔF, inds = _notrunc_ind(Δt);
             kwargs...
         )
-        foreachblock(Δt, t) do c, (Δb, b)
+        foreachblock(Δt) do c, (Δb,)
+            b = nothing_or_block(t, c)
             ind = get(inds, c, nothing)
             isnothing(ind) && return nothing
-            Fc = block.(F, Ref(c))
-            ΔFc = block.(ΔF, Ref(c))
+            Fc = nothing_or_block.(F, Ref(c))
+            ΔFc = nothing_or_block.(ΔF, Ref(c))
             MAK.$pullback!(Δb, b, Fc, ΔFc, ind; kwargs...)
             return nothing
         end
         return Δt
     end
+    @eval function MAK.$pullback!(
+            Δt::AbstractTensorMap, t, F, ΔF, ::Colon; kwargs...
+        )
+        return MAK.$pullback!(Δt, t, F, ΔF, _notrunc_ind(Δt); kwargs...)
+    end
 end
 
 for pullback_trunc! in (:svd_trunc_pullback!, :eig_trunc_pullback!, :eigh_trunc_pullback!)
     @eval function MAK.$pullback_trunc!(
-            Δt::AbstractTensorMap, t::AbstractTensorMap, F, ΔF; kwargs...
+            Δt::AbstractTensorMap, t, F, ΔF; kwargs...
         )
         foreachblock(Δt, t) do c, (Δb, b)
             Fc = block.(F, Ref(c))
@@ -97,3 +121,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

src/factorizations/pushforwards.jl

@@ -0,0 +1,57 @@
+for pushforward! in (
+        :qr_pushforward!, :lq_pushforward!, :left_polar_pushforward!, :right_polar_pushforward!,
+    )
+    @eval function MAK.$pushforward!(
+            Δt::AbstractTensorMap, t, F, ΔF; kwargs...
+        )
+        foreachblock(Δt) do c, (Δb,)
+            b = nothing_or_block(t, c)
+            Fc = block.(F, Ref(c))
+            ΔFc = block.(ΔF, Ref(c))
+            return MAK.$pushforward!(Δb, b, Fc, ΔFc; kwargs...)
+        end
+        return Δt
+    end
+end
+for pushforward! in (:qr_null_pushforward!, :lq_null_pushforward!)
+    @eval function MAK.$pushforward!(
+            Δt::AbstractTensorMap, t::AbstractTensorMap, F, ΔF; kwargs...
+        )
+        foreachblock(Δt, t) do c, (Δb, b)
+            Fc = block(F, c)
+            ΔFc = block(ΔF, c)
+            return MAK.$pushforward!(Δb, b, Fc, ΔFc; kwargs...)
+        end
+        return Δt
+    end
+end
+
+for pushforward! in (:eig_vals_pushforward!, :eigh_vals_pushforward!)
+    @eval function MAK.$pushforward!(
+            Δt::AbstractTensorMap, t, DV::Tuple{Diagonal, <:AbstractTensorMap}, ΔD, inds;
+            kwargs...
+        )
+        return MAK.$pushforward!(Δt, t, (MAK.diagonal(parent(DV[1])), DV[2]), ΔD, inds; kwargs...)
+    end
+end
+function MAK.svd_vals_pushforward!(
+        Δt::AbstractTensorMap, t, USVᴴ::Tuple{<:AbstractTensorMap, Diagonal, <:AbstractTensorMap}, ΔS, ind;
+        kwargs...
+    )
+    return MAK.svd_vals_pushforward!(Δt, t, (USVᴴ[1], MAK.diagonal(parent(USVᴴ[2])), USVᴴ[3]), ΔS, ind; kwargs...)
+end
+
+for pushforward! in (:svd_pushforward!, :eig_pushforward!, :eigh_pushforward!)
+    @eval function MAK.$pushforward!(
+            Δt::AbstractTensorMap, t, F, ΔF;
+            kwargs...
+        )
+        foreachblock(Δt, t) do c, (Δb, b)
+            Fc = nothing_or_block.(F, Ref(c))
+            ΔFc = nothing_or_block.(ΔF, Ref(c))
+            MAK.$pushforward!(Δb, b, Fc, ΔFc; kwargs...)
+            return nothing
+        end
+        return Δt
+    end
+end

test/enzyme/eig.jl

@@ -0,0 +1,37 @@
+using Test, TestExtras
+using TensorKit
+using TensorOperations
+using MatrixAlgebraKit
+using MatrixAlgebraKit: remove_eig_gauge_dependence!
+using Enzyme, EnzymeTestUtils
+using Random
+
+is_ci = get(ENV, "CI", "false") == "true"
+
+spacelist = ad_spacelist(fast_tests)
+eltypes = (Float64, ComplexF64)
+
+@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]))
+    atol = default_tol(T)
+    rtol = default_tol(T)
+
+    if !is_ci
+        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)
+        EnzymeTestUtils.test_forward(eig_full, Duplicated, (t, Duplicated); atol, rtol)
+
+        D = eig_vals(t)
+        EnzymeTestUtils.test_reverse(eig_vals, Duplicated, (t, Duplicated); atol, rtol)
+        EnzymeTestUtils.test_forward(eig_vals, Duplicated, (t, Duplicated); atol, rtol)
+    end
+
+    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

test/enzyme/eigh.jl

@@ -0,0 +1,39 @@
+using Test, TestExtras
+using TensorKit
+using TensorOperations
+using MatrixAlgebraKit
+using MatrixAlgebraKit: remove_eigh_gauge_dependence!
+using Enzyme, EnzymeTestUtils
+using Random
+
+is_ci = get(ENV, "CI", "false") == "true"
+
+spacelist = ad_spacelist(fast_tests)
+eltypes = (Float64, ComplexF64)
+
+@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]))
+    atol = default_tol(T)
+    rtol = default_tol(T)
+    th = project_hermitian(t)
+    if !is_ci
+        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)
+        EnzymeTestUtils.test_forward(proj_eigh_full, Duplicated, (th, Duplicated); atol, rtol)
+
+        D = eigh_vals(th)
+        EnzymeTestUtils.test_reverse(eigh_vals ∘ project_hermitian, Duplicated, (th, Duplicated); atol, rtol)
+        EnzymeTestUtils.test_forward(eigh_vals ∘ project_hermitian, Duplicated, (th, Duplicated); atol, rtol)
+    end
+
+    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

test/enzyme/lq.jl

@@ -0,0 +1,36 @@
+using Test, TestExtras
+using TensorKit
+using TensorOperations
+using MatrixAlgebraKit
+using MatrixAlgebraKit: remove_lq_gauge_dependence!, remove_lq_null_gauge_dependence!
+using Enzyme, EnzymeTestUtils
+using Random
+
+is_ci = get(ENV, "CI", "false") == "true"
+
+spacelist = ad_spacelist(fast_tests)
+eltypes = (Float64, ComplexF64)
+
+@timedtestset "Enzyme - Factorizations (LQ): $(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)
+
+    EnzymeTestUtils.test_reverse(lq_compact, Duplicated, (t, Duplicated); atol, rtol)
+    EnzymeTestUtils.test_forward(lq_compact, Duplicated, (t, Duplicated); atol, rtol)
+
+    if !is_ci
+        # 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)
+        EnzymeTestUtils.test_forward(lq_full, Duplicated, (t, Duplicated); 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)
+        EnzymeTestUtils.test_forward(lq_null, Duplicated, (t, Duplicated); atol, rtol)
+    end
+end

test/enzyme/projections.jl

@@ -0,0 +1,24 @@
+using Test, TestExtras
+using TensorKit
+using TensorOperations
+using MatrixAlgebraKit
+using Enzyme, EnzymeTestUtils
+using Random
+
+is_ci = get(ENV, "CI", "false") == "true"
+
+spacelist = ad_spacelist(fast_tests)
+eltypes = (Float64, ComplexF64)
+
+@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]))
+    atol = default_tol(T)
+    rtol = default_tol(T)
+    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)
+    EnzymeTestUtils.test_forward(project_hermitian, Duplicated, (t, Duplicated); atol, rtol)
+    EnzymeTestUtils.test_forward(project_antihermitian, Duplicated, (t, Duplicated); atol, rtol)
+    EnzymeTestUtils.test_forward(project_hermitian!, Duplicated, (t, Duplicated); atol, rtol)
+    EnzymeTestUtils.test_forward(project_antihermitian!, Duplicated, (t, Duplicated); atol, rtol)
+end

test/enzyme/qr.jl

@@ -0,0 +1,36 @@
+using Test, TestExtras
+using TensorKit
+using TensorOperations
+using MatrixAlgebraKit
+using MatrixAlgebraKit: remove_qr_gauge_dependence!, remove_qr_null_gauge_dependence!
+using Enzyme, EnzymeTestUtils
+using Random
+
+is_ci = get(ENV, "CI", "false") == "true"
+
+spacelist = ad_spacelist(fast_tests)
+eltypes = (Float64, ComplexF64)
+
+@timedtestset "Enzyme - Factorizations (QR): $(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)
+
+    EnzymeTestUtils.test_reverse(qr_compact, Duplicated, (t, Duplicated); atol, rtol)
+    EnzymeTestUtils.test_forward(qr_compact, Duplicated, (t, Duplicated); atol, rtol)
+
+    if !is_ci
+        # 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)
+        EnzymeTestUtils.test_forward(qr_full, Duplicated, (t, Duplicated); 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)
+        EnzymeTestUtils.test_forward(qr_null, Duplicated, (t, Duplicated); atol, rtol)
+    end
+end