[WIP] Get PEPSKit working with Enzyme

[kshyatt]

Trying to see if I can make this happen. Right now, SimultaneousCTMRG doesn't work because we need a reverse rule for dtmap (working on that) but other things... do seem to work? Kind of cool! There are some other errors I see related to LLVM I will open issues for as I encounter them. But, if others want to take a look, go for it :)

[github-actions]

Your PR requires formatting changes to meet the project's style guidelines.
Please consider running Runic (git runic main) to apply these changes.

Click here to view the suggested changes.

diff --git a/ext/PEPSKitEnzymeExt.jl b/ext/PEPSKitEnzymeExt.jl
index b0e4a75..0b71f77 100644
--- a/ext/PEPSKitEnzymeExt.jl
+++ b/ext/PEPSKitEnzymeExt.jl
@@ -12,11 +12,11 @@ using Enzyme.EnzymeCore: EnzymeRules
 @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}}
+        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)
@@ -41,12 +41,12 @@ function EnzymeRules.augmented_primal(
 end
 
 function EnzymeRules.reverse(
-    config::EnzymeRules.RevConfigWidth{1},
-    func::Const{typeof(MatrixAlgebraKit.svd_trunc)},
-    ::Type{RT},
-    cache,
-    t::Annotation,
-    alg::Const{<:SVDAdjoint{F, R}}
+        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⁺
@@ -57,17 +57,18 @@ function EnzymeRules.reverse(
             gauge_atol = gtol(dUSVᴴtrunc), degeneracy_atol = alg.val.rrule_alg.degeneracy_atol,
         )
     end
-    return ntuple(Returns(nothing), 3) 
+    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}
+        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)
@@ -75,26 +76,28 @@ function EnzymeRules.augmented_primal(
 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}
+        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) 
+    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}
+        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
@@ -107,13 +110,14 @@ function EnzymeRules.augmented_primal(
 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}
+        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)
diff --git a/test/ctmrg/pepo.jl b/test/ctmrg/pepo.jl
index b23afde..735fc5b 100644
--- a/test/ctmrg/pepo.jl
+++ b/test/ctmrg/pepo.jl
@@ -5,7 +5,7 @@ using PEPSKit
 using TensorKit
 using KrylovKit
 using OptimKit
-using Enzyme 
+using Enzyme
 
 ## Setup
 
diff --git a/test/gradients/ctmrg_gradients.jl b/test/gradients/ctmrg_gradients.jl
index e1b7127..20d839e 100644
--- a/test/gradients/ctmrg_gradients.jl
+++ b/test/gradients/ctmrg_gradients.jl
@@ -2,7 +2,7 @@ using Test
 using Random
 using PEPSKit
 using TensorKit
-using Enzyme 
+using Enzyme
 using OptimKit
 using KrylovKit
 
@@ -19,7 +19,7 @@ names = ["Heisenberg", "p-wave superconductor"]
 gradtol = 1.0e-4
 ctmrg_verbosity = 0
 #ctmrg_algs = [[:SequentialCTMRG, :SimultaneousCTMRG], [:SequentialCTMRG, :SimultaneousCTMRG]]
-ctmrg_algs = [[:SequentialCTMRG,], [:SequentialCTMRG,]]
+ctmrg_algs = [[:SequentialCTMRG], [:SequentialCTMRG]]
 projector_algs = [[:HalfInfiniteProjector, :FullInfiniteProjector], [:HalfInfiniteProjector, :FullInfiniteProjector]]
 svd_rrule_algs = [[:FullPullback, :TruncPullback, :Arnoldi], [:FullPullback, :Arnoldi]]
 #gradient_algs = [[nothing, :FixedPointGradient], [:FixedPointGradient]]
@@ -65,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(

[codecov]

Codecov Report

❌ Patch coverage is 0% with 54 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
ext/PEPSKitEnzymeExt.jl	0.00%	54 Missing ⚠️

Files with missing lines	Coverage Δ
ext/PEPSKitEnzymeExt.jl	0.00% <0.00%> (ø)

... and 38 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.