[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"
@@ -30,11 +31,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.196"
 FiniteDifferences = "0.12"
 KrylovKit = "0.9.5, 0.10"
 LinearAlgebra = "1"
@@ -54,3 +59,6 @@ TupleTools = "1.6.0"
 VectorInterface = "0.4, 0.5, 0.6"
 Zygote = "0.6, 0.7"
 julia = "1.10"
+
+[sources]
+Enzyme = {url = "https://github.com/EnzymeAD/Enzyme.jl/", rev = "ksh/icc"}

ext/PEPSKitEnzymeExt.jl

@@ -0,0 +1,182 @@
+module PEPSKitEnzymeExt
+
+using PEPSKit, MPSKit, TensorKit, MatrixAlgebraKit
+using PEPSKit: SVDAdjoint, EighAdjoint, QRAdjoint, CTMRGAlgorithm, FixedPointGradient, sdiag_pow, dtmap
+using ChainRulesCore: ignore_derivatives
+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
+
+@inline EnzymeRules.inactive(::typeof(PEPSKit.checklattice), args...) = nothing
+
+# Without this, Enzyme differentiates through `ignore_derivatives`
+@inline EnzymeRules.inactive(::typeof(ignore_derivatives), args...) = nothing
+
+function EnzymeRules.augmented_primal(
+        config::EnzymeRules.RevConfigWidth{1},
+        func::Const{typeof(MatrixAlgebraKit.svd_trunc_no_error)},
+        ::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)
+
+    output = (Ũ, S̃, Ṽ⁺, truncerror)
+    USVᴴtrunc = (Ũ, S̃, Ṽ⁺)
+    primal = EnzymeRules.needs_primal(config) ? USVᴴtrunc : nothing
+    # This creates new output shadow matrices, we use USVᴴtrunc 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_no_error)},
+        ::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⁺
+    gtol = PEPSKit._get_pullback_gauge_tol(alg.val.rrule_alg.verbosity)
+    if !isa(t, Const)
+        t.dval = MatrixAlgebraKit.svd_pullback!(
+            t.dval, t.val, (U, S, V⁺), dUSVᴴ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}
+    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
+    denv = isnothing(shadow) ? nothing : shadow[2]
+    primal = EnzymeRules.needs_primal(config) ? (env, info) : nothing
+    return EnzymeRules.AugmentedReturn(primal, shadow, (env_conv, env, denv, 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, env, denv, alg_gauge, alg_fixed = cache
+    signs, corner_phases, edge_phases = PEPSKit.compute_gauge_fix_gauge(env_conv, env, 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)}
+    env_vjp = Enzyme.autodiff_thunk(ReverseSplitWithPrimal, Duplicated, gauge_fixed_iteration, state, Duplicated(env, denv))
+    # 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(denv, ∂f∂x, ∂f∂A, denv, gradmode.solver_alg)
+    return ntuple(Returns(nothing), 4)
+end
+
+function EnzymeRules.augmented_primal(
+        config::EnzymeRules.RevConfigWidth{1},
+        ::Const{typeof(dtmap)},
+        ::Type{RT},
+        f::Const,
+        A::Annotation{<:AbstractArray},
+        scheduler::Annotation
+    ) where {RT}
+    if !isa(A, Const)
+        el_rrules = tmap(A.val, A.dval; scheduler.val) do a, da
+            Enzyme.autodiff_thunk(ReverseSplitWithPrimal, Duplicated, f, Duplicated(a, da))
+        end
+        y = map(first, el_rrules)
+    else
+        el_rrules = nothing
+        y = map(a -> a.val, A)
+    end
+    dy = map(Enzyme.make_zero, y)
+    shadow = EnzymeRules.needs_shadow(config) ? dy : nothing
+    primal = EnzymeRules.needs_primal(config) ? y : nothing
+    return EnzymeRules.AugmentedReturn(primal, shadow, (y, dy, el_rrules))
+end
+
+function EnzymeRules.reverse(
+        config::EnzymeRules.RevConfigWidth{1},
+        ::Const{typeof(dtmap)},
+        ::Type{RT},
+        cache,
+        f::Const,
+        A::Annotation{<:AbstractArray},
+        scheduler::Annotation
+    ) where {RT}
+    ys, dys, el_rrules = cache
+    backevals = tmap(el_rrules, dys; scheduler.val) do el_rrule, dy
+        last(el_rrule)(dy)
+    end
+    return (nothing, nothing, nothing)
+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"
@@ -26,6 +28,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

@@ -91,7 +91,7 @@ end
         dir = InfinitePEPS(Pspace, Vspace)
         psi = InfinitePEPS(Pspace, Vspace)
         # instantiate to avoid having to type this twice...
-        concrete_ctmrg_alg = PEPSKit.CTMRGAlgorithm(;
+        contrete_ctmrg_alg = PEPSKit.CTMRGAlgorithm(;
             alg = ctmrg_alg,
             verbosity = ctmrg_verbosity,
             projector_alg = projector_alg,
@@ -105,7 +105,7 @@ end
                 alg = gradient_alg, solver_alg = (; alg = gradient_solver_alg, tol = gradtol)
             )
         end
-        env, = leading_boundary(CTMRGEnv(psi, Espace), psi, concrete_ctmrg_alg)
+        env, = leading_boundary(CTMRGEnv(psi, Espace), psi, contrete_ctmrg_alg)
         alphas, fs, dfs1, dfs2 = OptimKit.optimtest(
             (psi, env),
             dir;
@@ -118,7 +118,7 @@ end
                     leading_boundary,
                     env,
                     psi,
-                    concrete_ctmrg_alg;
+                    contrete_ctmrg_alg;
                     alg_rrule = concrete_gradient_alg,
                 )
                 return cost_function(psi, env2, models[i])