From c437f21400c4bf49d687cee8efe192606e5d09a5 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Mon, 15 Jun 2026 12:18:26 +0200 Subject: [PATCH 1/2] Getting started with Enzyme --- Project.toml | 5 ++ ext/PEPSKitEnzymeExt.jl | 136 ++++++++++++++++++++++++++++++++++++++++ test/Project.toml | 3 + test/ctmrg/pepo.jl | 13 ++-- 4 files changed, 149 insertions(+), 8 deletions(-) create mode 100644 ext/PEPSKitEnzymeExt.jl diff --git a/Project.toml b/Project.toml index 68ece85cd..e67a9d99b 100644 --- a/Project.toml +++ b/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" diff --git a/ext/PEPSKitEnzymeExt.jl b/ext/PEPSKitEnzymeExt.jl new file mode 100644 index 000000000..1f60eb7e1 --- /dev/null +++ b/ext/PEPSKitEnzymeExt.jl @@ -0,0 +1,136 @@ +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) + println("IN AUGMENTED PRIMAL") + 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, = 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, info = PEPSKit.ctmrg_iteration(InfiniteSquareNetwork(state.val), env.val, alg_fixed) + shadow = Enzyne.make_zero((env, info)) + 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) + println("RUN AUTODIFF IN RRULE") + 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 diff --git a/test/Project.toml b/test/Project.toml index 2742a2812..22185e9d8 100644 --- a/test/Project.toml +++ b/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" diff --git a/test/ctmrg/pepo.jl b/test/ctmrg/pepo.jl index e2e0ec05a..b23afde4b 100644 --- a/test/ctmrg/pepo.jl +++ b/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 From bef7e910fd5e384b6a9d229547b034fbf0cb5d85 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Fri, 17 Jul 2026 16:58:57 +0200 Subject: [PATCH 2/2] Some gradient tests working --- ext/PEPSKitEnzymeExt.jl | 9 ++++--- test/gradients/ctmrg_gradients.jl | 40 +++++++++++++------------------ 2 files changed, 20 insertions(+), 29 deletions(-) diff --git a/ext/PEPSKitEnzymeExt.jl b/ext/PEPSKitEnzymeExt.jl index 1f60eb7e1..b0e4a75bd 100644 --- a/ext/PEPSKitEnzymeExt.jl +++ b/ext/PEPSKitEnzymeExt.jl @@ -69,7 +69,6 @@ function EnzymeRules.augmented_primal( f::Const, args::Annotation...) where {RT} alg_rrule = get(kw.val, :alg_rrule, nothing) - println("IN AUGMENTED PRIMAL") 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)) @@ -97,12 +96,13 @@ function EnzymeRules.augmented_primal( state::Annotation, alg::Const{<:CTMRGAlgorithm}) where {RT} #PEPSKit._check_algorithm_combination(alg, gradmode) - env, = MPSKit.leading_boundary(envinit.val, state.val, alg.val) + 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, info = PEPSKit.ctmrg_iteration(InfiniteSquareNetwork(state.val), env.val, alg_fixed) - shadow = Enzyne.make_zero((env, info)) + 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 @@ -123,7 +123,6 @@ function EnzymeRules.reverse( # prepare its pullback sig = Tuple{typeof(gauge_fixed_iteration), typeof(state), typeof(env)} rule = Mooncake.build_rrule(gauge_fixed_iteration, state, env) - println("RUN AUTODIFF IN RRULE") 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] diff --git a/test/gradients/ctmrg_gradients.jl b/test/gradients/ctmrg_gradients.jl index 8618e139a..e1b712786 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 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,6 +108,7 @@ end ) end env, = leading_boundary(CTMRGEnv(psi, Espace), psi, contrete_ctmrg_alg) + model = models[i] alphas, fs, dfs1, dfs2 = OptimKit.optimtest( (psi, env), dir; @@ -113,23 +116,17 @@ end 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=#