Conversation
Contributor
|
Your PR requires formatting changes to meet the project's style guidelines. 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( |
kshyatt
force-pushed
the
ksh/enzyme
branch
2 times, most recently
from
August 6, 2026 12:59
8a09a4d to
3a42eec
Compare
kshyatt
force-pushed
the
ksh/enzyme
branch
2 times, most recently
from
September 24, 2026 08:25
27e1312 to
891a8c8
Compare
Member
Author
|
As mentioned in a few comments, this is completely blocked on EnzymeAD/Enzyme#3152 which is needed for the |
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Trying to see if I can make this happen. Right now,
SimultaneousCTMRGdoesn't work because we need a reverse rule fordtmap(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 :)