[WIP] Get PEPSKit working with Enzyme

Project.toml

@@ -11,6 +11,7 @@ Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -29,11 +30,15 @@ TupleTools = "9d95972d-f1c8-5527-a6e0-b4b365fa01f6"
 VectorInterface = "409d34a3-91d5-4945-b6ec-7529ddf182d8"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
+[extensions]
+PEPSKitEnzymeExt = "Enzyme"
+
 [compat]
 Accessors = "0.1"
 ChainRulesCore = "1.0"
 Compat = "3.46, 4.2"
 DocStringExtensions = "0.9.3"
+Enzyme = "0.13.158"
 FiniteDifferences = "0.12"
 KrylovKit = "0.9.5, 0.10"
 LinearAlgebra = "1"

ext/PEPSKitEnzymeExt.jl

@@ -0,0 +1,135 @@
+module PEPSKitEnzymeExt
+
+using PEPSKit, MPSKit, TensorKit, MatrixAlgebraKit
+using PEPSKit: SVDAdjoint, EighAdjoint, QRAdjoint, CTMRGAlgorithm, FixedPointGradient, sdiag_pow
+import PEPSKit: real_inner
+using Enzyme
+using Enzyme.EnzymeCore: EnzymeRules
+
+@inline EnzymeRules.inactive_type(::Type{SVDAdjoint}) = true
+@inline EnzymeRules.inactive_type(::Type{QRAdjoint}) = true
+@inline EnzymeRules.inactive_type(::Type{EighAdjoint}) = true
+@inline EnzymeRules.inactive_type(::Type{CTMRGAlgorithm}) = true
+
+function EnzymeRules.augmented_primal(
+    config::EnzymeRules.RevConfigWidth{1},
+    func::Const{typeof(MatrixAlgebraKit.svd_trunc)},
+    ::Type{RT},
+    t::Annotation,
+    alg::Const{<:SVDAdjoint{F, R}}
+    ) where {RT, F, R <: PEPSKit.FullPullback}
+    # requires access to the full decomposition
+    U, S, V⁺ = svd_compact(t.val, alg.val.fwd_alg.alg)
+    (Ũ, S̃, Ṽ⁺), inds = MatrixAlgebraKit.truncate(svd_trunc!, (U, S, V⁺), alg.val.fwd_alg.trunc)
+    truncerror = MatrixAlgebraKit.truncation_error(diagview(S), inds)
+
+    gtol = PEPSKit._get_pullback_gauge_tol(alg.val.rrule_alg.verbosity)
+    output = (Ũ, S̃, Ṽ⁺, truncerror)
+    USVᴴtrunc = (Ũ, S̃, Ṽ⁺)
+    primal = EnzymeRules.needs_primal(config) ? USVᴴ′ : nothing
+    # This creates new output shadow matrices, we use USVᴴ′ to ensure the
+    # eltypes and dimensions are correct.
+    # These new shadow matrices are "filled in" with the accumulated
+    # results from earlier in reverse-mode AD after this function exits
+    # and before `reverse` is called.
+    dret = if EnzymeRules.needs_shadow(config)
+        (zero(USVᴴtrunc[1]), Diagonal(zero(USVᴴtrunc[2].diag)), zero(USVᴴtrunc[3]))
+    else
+        nothing
+    end
+    return EnzymeRules.AugmentedReturn(primal, dret, (dret, (U, S, V⁺), inds))
+end
+
+function EnzymeRules.reverse(
+    config::EnzymeRules.RevConfigWidth{1},
+    func::Const{typeof(MatrixAlgebraKit.svd_trunc)},
+    ::Type{RT},
+    cache,
+    t::Annotation,
+    alg::Const{<:SVDAdjoint{F, R}}
+    ) where {RT, F, R <: PEPSKit.FullPullback}
+    dUSVᴴtrunc, USV⁺, ind = cache
+    U, S, V⁺ = USV⁺
+    _warn_pullback_truncerror(dϵ)
+    if !isa(t, Const)
+        t.dval = MatrixAlgebraKit.svd_pullback!(
+            t.dval, t.val, (U, S, V⁺), ΔUSVᴴtrunc, ind;
+            gauge_atol = gtol(dUSVᴴtrunc), degeneracy_atol = alg.val.rrule_alg.degeneracy_atol,
+        )
+    end
+    return ntuple(Returns(nothing), 3) 
+end
+
+function EnzymeRules.augmented_primal(
+    config::EnzymeRules.RevConfigWidth{1},
+    ::Const{typeof(Core.kwcall)},
+    ::Type{RT},
+    kw::Const{<:NamedTuple},
+    ::Const{typeof(PEPSKit.hook_pullback)},
+    f::Const,
+    args::Annotation...) where {RT}
+    alg_rrule = get(kw.val, :alg_rrule, nothing)
+    primal, rrule_func = PEPSKit._rrule(alg_rrule, f.val, map(arg -> getfield(arg, :val), args)...)
+    shadow = Enzyme.make_zero(primal)
+    return EnzymeRules.AugmentedReturn(primal, shadow, (shadow, rrule_func))
+end
+
+function EnzymeRules.reverse(
+    config::EnzymeRules.RevConfigWidth{1},
+    ::Const{typeof(Core.kwcall)},
+    ::Type{RT},
+    cache,
+    kw::Const{<:NamedTuple},
+    ::Const{typeof(PEPSKit.hook_pullback)},
+    args::Annotation...) where {RT}
+    println("IN REVERSE")
+    shadow, rrule_func = cache
+    rrule_func(shadow) 
+    return ntuple(Returns(nothing), 2 + length(args))
+end
+
+function EnzymeRules.augmented_primal(
+    config::EnzymeRules.RevConfigWidth{1},
+    ::Const{typeof(MPSKit.leading_boundary)},
+    ::Type{RT},
+    envinit::Annotation,
+    state::Annotation,
+    alg::Const{<:CTMRGAlgorithm}) where {RT}
+    #PEPSKit._check_algorithm_combination(alg, gradmode)
+    env, info = MPSKit.leading_boundary(envinit.val, state.val, alg.val)
+    # prepare iterating function corresponding to a single gauge-fixed CTMRG iteration
+    alg_fixed = PEPSKit._set_fixed_truncation(alg.val) # fix spaces during differentiation
+    alg_gauge = PEPSKit._scrambling_env_gauge(alg.val) # select appropriate gauge-fixing algorithm
+    env_conv, _ = PEPSKit.ctmrg_iteration(InfiniteSquareNetwork(state.val), env, alg_fixed)
+    shadow = EnzymeRules.needs_shadow(config) ? Enzyme.make_zero((env, info)) : nothing
+    primal = EnzymeRules.needs_primal(config) ? (env, info) : nothing
+    return EnzymeRules.AugmentedReturn(primal, shadow, (env_conv, alg_gauge, alg_fixed))
+end
+
+function EnzymeRules.reverse(
+    config::EnzymeRules.RevConfigWidth{1},
+    ::Const{typeof(MPSKit.leading_boundary)},
+    ::Type{RT},
+    cache,
+    envinit::Annotation,
+    state::Annotation,
+    alg::Const{<:CTMRGAlgorithm}) where {RT}
+    env_conv, alg_gauge, alg_fixed = cache
+    signs, corner_phases, edge_phases = PEPSKit.compute_gauge_fix_gauge(env_conv, env.val, alg_gauge)
+    function gauge_fixed_iteration(A, x)
+        x′ = PEPSKit.ctmrg_iteration(InfiniteSquareNetwork(A), x, alg_fixed)[1]
+        return PEPSKit.fix_phases(x′, signs, corner_phases, edge_phases)
+    end
+    # prepare its pullback
+    sig = Tuple{typeof(gauge_fixed_iteration), typeof(state), typeof(env)}
+    rule = Mooncake.build_rrule(gauge_fixed_iteration, state, env)
+    env_vjp = Enzyme.autodiff(rule, gauge_fixed_iteration, state, env)
+    # split off state and environment parts
+    ∂f∂A(x)::typeof(state) = env_vjp(x)[2]
+    ∂f∂x(x)::typeof(env) = env_vjp(x)[3]
+    # evaluate the geometric sum
+    PEPSKit.fixedpoint_gradient(env.dval, ∂f∂x, ∂f∂A, env.dval, gradmode.solver_alg)
+    return ntuple(Returns(NoRData()), 4)
+end
+
+end

test/Project.toml

@@ -4,6 +4,8 @@ name = "PEPSKitTests"
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+EnzymeTestUtils = "12d8515a-0907-448a-8884-5fe00fdf1c5a"
 KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MPSKit = "bb1c41ca-d63c-52ed-829e-0820dda26502"
@@ -25,6 +27,7 @@ PEPSKit = {path = ".."}
 
 [compat]
 ChainRulesTestUtils = "1.13"
+EnzymeTestUtils = "0.2.8"
 ParallelTestRunner = "2.6.0"
 QuadGK = "2.11.1"
 Test = "1"

test/ctmrg/pepo.jl

@@ -5,7 +5,7 @@ using PEPSKit
 using TensorKit
 using KrylovKit
 using OptimKit
-using Zygote
+using Enzyme 
 
 ## Setup
 
@@ -112,15 +112,11 @@ end
         retract = pepo_retract,
         (transport!) = (pepo_transport!),
     ) do (psi, env2, env3)
-        E, gs = withgradient(psi) do ψ
+        function energ(ψ)
             n2 = InfiniteSquareNetwork(ψ)
-            env2′, info = PEPSKit.hook_pullback(
-                leading_boundary, env2, n2, ctm_alg; alg_rrule = gradient_alg
-            )
+            env2′, info = leading_boundary(env2, n2, ctm_alg)
             n3 = InfiniteSquareNetwork(ψ, T)
-            env3′, info = PEPSKit.hook_pullback(
-                leading_boundary, env3, n3, ctm_alg; alg_rrule = gradient_alg
-            )
+            env3′, info = leading_boundary(env3, n3, ctm_alg)
             PEPSKit.ignore_derivatives() do
                 PEPSKit.update!(env2, env2′)
                 PEPSKit.update!(env3, env3′)
@@ -129,6 +125,7 @@ end
             λ2 = network_value(n2, env2)
             return -log(real(λ3 / λ2))
         end
+        E, gs = Enzyme.autodiff(ReverseWithPrimal, Const(energ), Active, Duplicated(psi, zerovector(psi)))
         g = only(gs)
         return E, g
     end

test/gradients/ctmrg_gradients.jl

@@ -2,7 +2,7 @@ using Test
 using Random
 using PEPSKit
 using TensorKit
-using Zygote
+using Enzyme 
 using OptimKit
 using KrylovKit
 
@@ -18,10 +18,12 @@ names = ["Heisenberg", "p-wave superconductor"]
 
 gradtol = 1.0e-4
 ctmrg_verbosity = 0
-ctmrg_algs = [[:SequentialCTMRG, :SimultaneousCTMRG], [:SequentialCTMRG, :SimultaneousCTMRG]]
+#ctmrg_algs = [[:SequentialCTMRG, :SimultaneousCTMRG], [:SequentialCTMRG, :SimultaneousCTMRG]]
+ctmrg_algs = [[:SequentialCTMRG,], [:SequentialCTMRG,]]
 projector_algs = [[:HalfInfiniteProjector, :FullInfiniteProjector], [:HalfInfiniteProjector, :FullInfiniteProjector]]
 svd_rrule_algs = [[:FullPullback, :TruncPullback, :Arnoldi], [:FullPullback, :Arnoldi]]
-gradient_algs = [[nothing, :FixedPointGradient], [:FixedPointGradient]]
+#gradient_algs = [[nothing, :FixedPointGradient], [:FixedPointGradient]]
+gradient_algs = [[:FixedPointGradient], [:FixedPointGradient]]
 gradient_solver_algs = [
     [:GeomSum, :ManualIter, :GMRES, :BiCGStab, :Arnoldi],
     [:GeomSum, :ManualIter, :GMRES, :BiCGStab, :Arnoldi],
@@ -63,7 +65,7 @@ end
             ctmrg_alg, projector_alg, svd_rrule_alg, gradient_alg, gradient_solver_alg,
         ) in Iterators.product(
             calgs, palgs, salgs, galgs, gsalgs
-        )
+           )
 
         # filter disallowed algorithm combinations
         if _check_disallowed_combination(
@@ -106,30 +108,25 @@ end
             )
         end
         env, = leading_boundary(CTMRGEnv(psi, Espace), psi, contrete_ctmrg_alg)
+        model = models[i]
         alphas, fs, dfs1, dfs2 = OptimKit.optimtest(
             (psi, env),
             dir;
             alpha = steps,
             retract = PEPSKit.peps_retract,
             inner = PEPSKit.real_inner,
         ) do (peps, env)
-            E, g = Zygote.withgradient(peps) do psi
-                env2, = PEPSKit.hook_pullback(
-                    leading_boundary,
-                    env,
-                    psi,
-                    contrete_ctmrg_alg;
-                    alg_rrule = concrete_gradient_alg,
-                )
-                return cost_function(psi, env2, models[i])
+            function energ(ψ)
+                env2, info = leading_boundary(env, ψ, contrete_ctmrg_alg)
+                cost_function(ψ, env2, model)
             end
-
+            E, gs = Enzyme.autodiff(ReverseWithPrimal, Const(energ), Active, Duplicated(peps, zerovector(peps)))
             return E, only(g)
         end
         @test dfs1 ≈ dfs2 atol = 1.0e-2
     end
 end
-
+#=
 ## Regression test for gradient accuracy (https://github.com/QuantumKitHub/PEPSKit.jl/pull/276)
 @testset "AD CTMRG energy gradient accuracy regression test (#276)" begin
     Random.seed!(1234)
@@ -138,16 +135,11 @@ end
     gradient_alg = PEPSKit.GradientAlgorithm(; tol = 5.0e-8)
 
     function fg((peps, env))
-        E, g = Zygote.withgradient(peps) do ψ
-            env2, = PEPSKit.hook_pullback(
-                leading_boundary,
-                env,
-                ψ,
-                boundary_alg;
-                alg_rrule = gradient_alg,
-            )
+        function energ(ψ)
+            env2, = leading_boundary(env, ψ, boundary_alg)
             return cost_function(ψ, env2, H)
         end
+        E, gs = Enzyme.autodiff(ReverseWithPrimal, Const(energ), Active, Duplicated(peps, zerovector(peps)))
         return E, only(g)
     end
 
@@ -167,4 +159,4 @@ end
 
     # verify high gradient accuracy for small finite-difference step size
     @test dfs1 ≈ dfs2 rtol = 1.0e-2 * Δx
-end
+end=#