From cf33c358c529fbf36f87c766181631b6f8547940 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Mon, 15 Jun 2026 12:18:26 +0200 Subject: [PATCH 1/9] 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 77f66c9ce..1d8b88656 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" @@ -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.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 1961b9868..df2810366 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" @@ -26,6 +28,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 f67ce608c23fcda7a654369beb5030386be2cd13 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Fri, 17 Jul 2026 16:58:57 +0200 Subject: [PATCH 2/9] Some gradient tests working --- ext/PEPSKitEnzymeExt.jl | 9 ++++----- test/gradients/ctmrg_gradients.jl | 6 +++--- 2 files changed, 7 insertions(+), 8 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 09e3f47c4..8618e139a 100644 --- a/test/gradients/ctmrg_gradients.jl +++ b/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]) From c8f4d7fa837fefdad1bd0e41c6fb58ba485ac4de Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Mon, 27 Jul 2026 16:00:33 +0200 Subject: [PATCH 3/9] Some more Enzyme fixes --- ext/PEPSKitEnzymeExt.jl | 42 +++++++++++++++++++++++++++++++++++++---- 1 file changed, 38 insertions(+), 4 deletions(-) diff --git a/ext/PEPSKitEnzymeExt.jl b/ext/PEPSKitEnzymeExt.jl index b0e4a75bd..ab5d13fa9 100644 --- a/ext/PEPSKitEnzymeExt.jl +++ b/ext/PEPSKitEnzymeExt.jl @@ -1,7 +1,7 @@ module PEPSKitEnzymeExt using PEPSKit, MPSKit, TensorKit, MatrixAlgebraKit -using PEPSKit: SVDAdjoint, EighAdjoint, QRAdjoint, CTMRGAlgorithm, FixedPointGradient, sdiag_pow +using PEPSKit: SVDAdjoint, EighAdjoint, QRAdjoint, CTMRGAlgorithm, FixedPointGradient, sdiag_pow, dtmap import PEPSKit: real_inner using Enzyme using Enzyme.EnzymeCore: EnzymeRules @@ -11,6 +11,8 @@ using Enzyme.EnzymeCore: EnzymeRules @inline EnzymeRules.inactive_type(::Type{EighAdjoint}) = true @inline EnzymeRules.inactive_type(::Type{CTMRGAlgorithm}) = true +@inline EnzymeRules.inactive(::typeof(PEPSKit.checklattice), args...) = nothing + function EnzymeRules.augmented_primal( config::EnzymeRules.RevConfigWidth{1}, func::Const{typeof(MatrixAlgebraKit.svd_trunc)}, @@ -82,7 +84,6 @@ function EnzymeRules.reverse( 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)) @@ -122,8 +123,7 @@ function EnzymeRules.reverse( 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) + env_vjp = Enzyme.autodiff_thunk(ReverseWithPrimal, Duplicated, 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] @@ -132,4 +132,38 @@ function EnzymeRules.reverse( return ntuple(Returns(NoRData()), 4) end +function EnzymeRules.augmented_primal( + config::EnzymeRules.RevConfigWidth{1}, + ::Const{typeof(dtmap)}, + ::Type{RT}, + f::Const, + A::Annotation{<:AbstractArray}, + scheduler::Annotation + ) where {RT} + el_rrules = tmap(A.val, A.dval; scheduler.val) do a, da + Enzyme.autodiff_thunk(ReverseWithPrimal, Duplicated, f, Duplicated(a, da)) + end + y = map(first, el_rrules) + 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 From 791ff26b180a760a07251d84cab9383222a34f01 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Mon, 3 Aug 2026 11:01:48 +0200 Subject: [PATCH 4/9] More incremental --- Project.toml | 2 +- src/algorithms/toolbox.jl | 6 +- test/gradients/enz_ctmrg_gradients.jl | 164 ++++++++++++++++++++++++++ 3 files changed, 169 insertions(+), 3 deletions(-) create mode 100644 test/gradients/enz_ctmrg_gradients.jl diff --git a/Project.toml b/Project.toml index 1d8b88656..872296b56 100644 --- a/Project.toml +++ b/Project.toml @@ -39,7 +39,7 @@ Accessors = "0.1" ChainRulesCore = "1.0" Compat = "3.46, 4.2" DocStringExtensions = "0.9.3" -Enzyme = "0.13.158" +Enzyme = "0.13.196" FiniteDifferences = "0.12" KrylovKit = "0.9.5, 0.10" LinearAlgebra = "1" diff --git a/src/algorithms/toolbox.jl b/src/algorithms/toolbox.jl index b53e0f443..b8f28c78d 100644 --- a/src/algorithms/toolbox.jl +++ b/src/algorithms/toolbox.jl @@ -11,7 +11,8 @@ function MPSKit.expectation_value( bra::S, O::LocalOperator, ket::S, env::CTMRGEnv ) where {S <: InfiniteState} checklattice(bra, O, ket) - term_vals = dtmap(collect(O.terms)) do (inds, operator) # OhMyThreads can't iterate over O.terms directly + # TODO RESTORE THIS TO DTMAP + term_vals = map(collect(O.terms)) do (inds, operator) # OhMyThreads can't iterate over O.terms directly ρ = reduced_densitymatrix(inds, ket, bra, env) return trmul(operator, ρ) end @@ -22,7 +23,8 @@ function MPSKit.expectation_value( state::InfinitePEPO, O::LocalOperator, env::CTMRGEnv ) checklattice(state, O) - term_vals = dtmap(collect(O.terms)) do (inds, operator) # OhMyThreads can't iterate over O.terms directly + # TODO RESTORE THIS TO DTMAP + term_vals = map(collect(O.terms)) do (inds, operator) # OhMyThreads can't iterate over O.terms directly ρ = reduced_densitymatrix(inds, state, env) return trmul(operator, ρ) end diff --git a/test/gradients/enz_ctmrg_gradients.jl b/test/gradients/enz_ctmrg_gradients.jl new file mode 100644 index 000000000..96661e41b --- /dev/null +++ b/test/gradients/enz_ctmrg_gradients.jl @@ -0,0 +1,164 @@ +using Test +using Random +using PEPSKit +using TensorKit +using Enzyme +using OptimKit +using KrylovKit + +Enzyme.Compiler.VERBOSE_ERRORS[] = true + +## Test models, gradmodes and CTMRG algorithm +# ------------------------------------------- +χbond = 2 +χenv = 6 +Pspaces = [ComplexSpace(2), Vect[FermionParity](0 => 1, 1 => 1)] +Vspaces = [ComplexSpace(χbond), Vect[FermionParity](0 => χbond / 2, 1 => χbond / 2)] +Espaces = [ComplexSpace(χenv), Vect[FermionParity](0 => χenv / 2, 1 => χenv / 2)] +models = [heisenberg_XYZ(InfiniteSquare()), pwave_superconductor(InfiniteSquare())] +names = ["Heisenberg", "p-wave superconductor"] + +gradtol = 1.0e-4 +ctmrg_verbosity = 0 +#ctmrg_algs = [[:SequentialCTMRG, :SimultaneousCTMRG], [:SequentialCTMRG, :SimultaneousCTMRG]] +ctmrg_algs = [[:SimultaneousCTMRG], [:SimultaneousCTMRG]] +projector_algs = [[:HalfInfiniteProjector, :FullInfiniteProjector], [:HalfInfiniteProjector, :FullInfiniteProjector]] +svd_rrule_algs = [[:FullPullback, :TruncPullback, :Arnoldi], [:FullPullback, :Arnoldi]] +#gradient_algs = [[nothing, :FixedPointGradient], [:FixedPointGradient]] +gradient_algs = [[:FixedPointGradient], [:FixedPointGradient]] +gradient_solver_algs = [ + [:GeomSum, :ManualIter, :GMRES, :BiCGStab, :Arnoldi], + [:GeomSum, :ManualIter, :GMRES, :BiCGStab, :Arnoldi], +] +steps = -0.01:0.005:0.01 + +# don't check naive AD gradients for all algorithm combinations, since it's slow +naive_gradient_combinations = [ + (:SimultaneousCTMRG, :HalfInfiniteProjector, :FullPullback), + (:SimultaneousCTMRG, :FullInfiniteProjector, :FullPullback), + (:SequentialCTMRG, :HalfInfiniteProjector, :FullPullback), +] +naive_gradient_done = Set() + +# fixed-point differentiation is incompatible with sequential CTMRG +function _check_disallowed_combination( + ctmrg_alg, projector_alg, decomposition_rrule_alg, gradient_alg + ) + ctmrg_alg == :SequentialCTMRG && !isnothing(gradient_alg) && return true + return false +end + + +## Tests +# ------ +@testset "AD CTMRG energy gradients for $(names[i]) model" verbose = true for i in + eachindex( + models + ) + Pspace = Pspaces[i] + Vspace = Vspaces[i] + Espace = Espaces[i] + calgs = ctmrg_algs[i] + palgs = projector_algs[i] + salgs = svd_rrule_algs[i] + galgs = gradient_algs[i] + gsalgs = gradient_solver_algs[i] + @testset "ctmrg_alg=:$ctmrg_alg, projector_alg=:$projector_alg, svd_rrule_alg=:$svd_rrule_alg, gradient_alg=(; alg = :$gradient_alg, solver_alg = (; alg = :$gradient_solver_alg))" for ( + 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( + ctmrg_alg, projector_alg, svd_rrule_alg, gradient_alg + ) + # but verify that its use would throw an error + @test_throws ArgumentError PEPSOptimize(; + boundary_alg = (; alg = ctmrg_alg, projector_alg, decomposition_alg = (; rrule_alg = (; alg = svd_rrule_alg))), + gradient_alg = (; alg = gradient_alg, solver_alg = (; alg = gradient_solver_alg, tol = gradtol)), + ) + continue + end + + # check for allowed algorithm combinations when testing naive gradient + if isnothing(gradient_alg) + combo = (ctmrg_alg, projector_alg, svd_rrule_alg) + combo in naive_gradient_combinations || continue + combo in naive_gradient_done && continue + push!(naive_gradient_done, combo) + gradient_solver_alg = nothing # unused in naive gradient, so set to nothing to avoid confusion + end + + @info "optimtest of ctmrg_alg=:$ctmrg_alg, projector_alg=:$projector_alg, svd_rrule_alg=:$svd_rrule_alg and gradient_alg=(; alg = :$gradient_alg, solver_alg = (; alg = :$gradient_solver_alg)) on $(names[i])" + Random.seed!(42039482030) + dir = InfinitePEPS(Pspace, Vspace) + psi = InfinitePEPS(Pspace, Vspace) + # instantiate to avoid having to type this twice... + contrete_ctmrg_alg = PEPSKit.CTMRGAlgorithm(; + alg = ctmrg_alg, + verbosity = ctmrg_verbosity, + projector_alg = projector_alg, + decomposition_alg = SVDAdjoint(; rrule_alg = (; alg = svd_rrule_alg)), + ) + # instantiate because hook_pullback doesn't go through the keyword selector... + concrete_gradient_alg = if isnothing(gradient_alg) + nothing # TODO: add this to the PEPSKit.GradientAlgorithm selector? + else + PEPSKit.GradientAlgorithm(; + alg = gradient_alg, solver_alg = (; alg = gradient_solver_alg, tol = gradtol) + ) + 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) + function energ(psi) + env2, info = leading_boundary(env, psi, contrete_ctmrg_alg) + cost_function(psi, env2, model) + end + E, gs = Enzyme.autodiff(set_runtime_activity(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) + + boundary_alg = PEPSKit.CTMRGAlgorithm(; tol = 1.0e-10) + gradient_alg = PEPSKit.GradientAlgorithm(; tol = 5.0e-8) + + function fg((peps, env)) + 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 + + # initialize randomly + H = heisenberg_XYZ(InfiniteSquare(1, 1)) + peps = PEPSKit.peps_normalize(InfinitePEPS(randn, ComplexF64, physicalspace(H)[1], ComplexSpace(3))) + env0 = CTMRGEnv(randn, ComplexF64, peps, ComplexSpace(20)) + + # test gradient against finite-difference + Δx = 1.0e-5 + _, _, dfs1, dfs2 = OptimKit.optimtest( + fg, (peps, env0); + alpha = LinRange(-Δx, Δx, 2), + retract = PEPSKit.peps_retract, + inner = PEPSKit.real_inner, + ) + + # verify high gradient accuracy for small finite-difference step size + @test dfs1 ≈ dfs2 rtol = 1.0e-2 * Δx +end=# From f9417311dcd17c8fdb43f7a62bcde8e9eaba2722 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Wed, 5 Aug 2026 12:30:29 +0200 Subject: [PATCH 5/9] Some fixups for rules --- ext/PEPSKitEnzymeExt.jl | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/ext/PEPSKitEnzymeExt.jl b/ext/PEPSKitEnzymeExt.jl index ab5d13fa9..6cd3a1a2b 100644 --- a/ext/PEPSKitEnzymeExt.jl +++ b/ext/PEPSKitEnzymeExt.jl @@ -25,11 +25,10 @@ function EnzymeRules.augmented_primal( (Ũ, 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 + 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 @@ -52,7 +51,8 @@ function EnzymeRules.reverse( ) where {RT, F, R <: PEPSKit.FullPullback} dUSVᴴtrunc, USV⁺, ind = cache U, S, V⁺ = USV⁺ - _warn_pullback_truncerror(dϵ) + MatrixAlgebraKit._warn_pullback_truncerror(dϵ) + 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⁺), ΔUSVᴴtrunc, ind; @@ -104,7 +104,7 @@ function EnzymeRules.augmented_primal( 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)) + return EnzymeRules.AugmentedReturn(primal, shadow, (env_conv, env, alg_gauge, alg_fixed)) end function EnzymeRules.reverse( @@ -115,7 +115,7 @@ function EnzymeRules.reverse( envinit::Annotation, state::Annotation, alg::Const{<:CTMRGAlgorithm}) where {RT} - env_conv, alg_gauge, alg_fixed = cache + env_conv, env, 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] From 864f670ea4f3d962f2600438fd9b89e9c8a38895 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Thu, 6 Aug 2026 10:05:13 +0200 Subject: [PATCH 6/9] Some more incremental fixes --- ext/PEPSKitEnzymeExt.jl | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/ext/PEPSKitEnzymeExt.jl b/ext/PEPSKitEnzymeExt.jl index 6cd3a1a2b..d4ccfe3c6 100644 --- a/ext/PEPSKitEnzymeExt.jl +++ b/ext/PEPSKitEnzymeExt.jl @@ -115,7 +115,7 @@ function EnzymeRules.reverse( envinit::Annotation, state::Annotation, alg::Const{<:CTMRGAlgorithm}) where {RT} - env_conv, env, alg_gauge, alg_fixed = cache + env_conv, env, denv, 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] @@ -123,13 +123,13 @@ function EnzymeRules.reverse( end # prepare its pullback sig = Tuple{typeof(gauge_fixed_iteration), typeof(state), typeof(env)} - env_vjp = Enzyme.autodiff_thunk(ReverseWithPrimal, Duplicated, gauge_fixed_iteration, state, 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(env.dval, ∂f∂x, ∂f∂A, env.dval, gradmode.solver_alg) - return ntuple(Returns(NoRData()), 4) + PEPSKit.fixedpoint_gradient(denv, ∂f∂x, ∂f∂A, denv, gradmode.solver_alg) + return ntuple(Returns(nothing), 4) end function EnzymeRules.augmented_primal( @@ -140,10 +140,14 @@ function EnzymeRules.augmented_primal( A::Annotation{<:AbstractArray}, scheduler::Annotation ) where {RT} - el_rrules = tmap(A.val, A.dval; scheduler.val) do a, da - Enzyme.autodiff_thunk(ReverseWithPrimal, Duplicated, f, Duplicated(a, da)) + 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 + y = map(a->a.val, A) end - y = map(first, el_rrules) dy = map(Enzyme.make_zero, y) shadow = EnzymeRules.needs_shadow(config) ? dy : nothing primal = EnzymeRules.needs_primal(config) ? y : nothing From 3a42eec29cc17ee1b837c2b35f978ddc094b4834 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Thu, 6 Aug 2026 14:59:18 +0200 Subject: [PATCH 7/9] Formatter --- ext/PEPSKitEnzymeExt.jl | 114 +++++++++++++------------- test/ctmrg/pepo.jl | 2 +- test/gradients/enz_ctmrg_gradients.jl | 4 +- 3 files changed, 62 insertions(+), 58 deletions(-) diff --git a/ext/PEPSKitEnzymeExt.jl b/ext/PEPSKitEnzymeExt.jl index d4ccfe3c6..55e64c376 100644 --- a/ext/PEPSKitEnzymeExt.jl +++ b/ext/PEPSKitEnzymeExt.jl @@ -14,11 +14,11 @@ using Enzyme.EnzymeCore: EnzymeRules @inline EnzymeRules.inactive(::typeof(PEPSKit.checklattice), args...) = nothing 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) @@ -42,12 +42,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⁺ @@ -59,17 +59,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) @@ -77,25 +78,27 @@ 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} 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 @@ -108,13 +111,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, env, denv, 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) @@ -133,12 +137,12 @@ function EnzymeRules.reverse( end function EnzymeRules.augmented_primal( - config::EnzymeRules.RevConfigWidth{1}, - ::Const{typeof(dtmap)}, - ::Type{RT}, - f::Const, - A::Annotation{<:AbstractArray}, - scheduler::Annotation + 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 @@ -146,7 +150,7 @@ function EnzymeRules.augmented_primal( end y = map(first, el_rrules) else - y = map(a->a.val, A) + y = map(a -> a.val, A) end dy = map(Enzyme.make_zero, y) shadow = EnzymeRules.needs_shadow(config) ? dy : nothing @@ -155,19 +159,19 @@ function EnzymeRules.augmented_primal( end function EnzymeRules.reverse( - config::EnzymeRules.RevConfigWidth{1}, - ::Const{typeof(dtmap)}, - ::Type{RT}, - cache, - f::Const, - A::Annotation{<:AbstractArray}, - scheduler::Annotation + 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) + return (nothing, nothing, nothing) end end diff --git a/test/ctmrg/pepo.jl b/test/ctmrg/pepo.jl index b23afde4b..735fc5bdf 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/enz_ctmrg_gradients.jl b/test/gradients/enz_ctmrg_gradients.jl index 96661e41b..a908cf475 100644 --- a/test/gradients/enz_ctmrg_gradients.jl +++ b/test/gradients/enz_ctmrg_gradients.jl @@ -2,7 +2,7 @@ using Test using Random using PEPSKit using TensorKit -using Enzyme +using Enzyme using OptimKit using KrylovKit @@ -67,7 +67,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( From 360a8686eff86fcf9775faf76dda9bec0b614040 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Thu, 6 Aug 2026 15:03:49 +0200 Subject: [PATCH 8/9] Restore dtmap --- src/algorithms/toolbox.jl | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/algorithms/toolbox.jl b/src/algorithms/toolbox.jl index b8f28c78d..b53e0f443 100644 --- a/src/algorithms/toolbox.jl +++ b/src/algorithms/toolbox.jl @@ -11,8 +11,7 @@ function MPSKit.expectation_value( bra::S, O::LocalOperator, ket::S, env::CTMRGEnv ) where {S <: InfiniteState} checklattice(bra, O, ket) - # TODO RESTORE THIS TO DTMAP - term_vals = map(collect(O.terms)) do (inds, operator) # OhMyThreads can't iterate over O.terms directly + term_vals = dtmap(collect(O.terms)) do (inds, operator) # OhMyThreads can't iterate over O.terms directly ρ = reduced_densitymatrix(inds, ket, bra, env) return trmul(operator, ρ) end @@ -23,8 +22,7 @@ function MPSKit.expectation_value( state::InfinitePEPO, O::LocalOperator, env::CTMRGEnv ) checklattice(state, O) - # TODO RESTORE THIS TO DTMAP - term_vals = map(collect(O.terms)) do (inds, operator) # OhMyThreads can't iterate over O.terms directly + term_vals = dtmap(collect(O.terms)) do (inds, operator) # OhMyThreads can't iterate over O.terms directly ρ = reduced_densitymatrix(inds, state, env) return trmul(operator, ρ) end From 7b50515afcf66ab0c3620a19c2aeee684d57292f Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Tue, 11 Aug 2026 14:16:23 +0200 Subject: [PATCH 9/9] Touchups --- Project.toml | 3 +++ ext/PEPSKitEnzymeExt.jl | 17 +++++++++++------ 2 files changed, 14 insertions(+), 6 deletions(-) diff --git a/Project.toml b/Project.toml index 872296b56..0984b4c79 100644 --- a/Project.toml +++ b/Project.toml @@ -59,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"} diff --git a/ext/PEPSKitEnzymeExt.jl b/ext/PEPSKitEnzymeExt.jl index 55e64c376..07060c2ab 100644 --- a/ext/PEPSKitEnzymeExt.jl +++ b/ext/PEPSKitEnzymeExt.jl @@ -2,6 +2,7 @@ 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 @@ -13,9 +14,12 @@ using Enzyme.EnzymeCore: EnzymeRules @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)}, + func::Const{typeof(MatrixAlgebraKit.svd_trunc_no_error)}, ::Type{RT}, t::Annotation, alg::Const{<:SVDAdjoint{F, R}} @@ -43,7 +47,7 @@ end function EnzymeRules.reverse( config::EnzymeRules.RevConfigWidth{1}, - func::Const{typeof(MatrixAlgebraKit.svd_trunc)}, + func::Const{typeof(MatrixAlgebraKit.svd_trunc_no_error)}, ::Type{RT}, cache, t::Annotation, @@ -51,11 +55,10 @@ function EnzymeRules.reverse( ) where {RT, F, R <: PEPSKit.FullPullback} dUSVᴴtrunc, USV⁺, ind = cache U, S, V⁺ = USV⁺ - MatrixAlgebraKit._warn_pullback_truncerror(dϵ) 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⁺), ΔUSVᴴtrunc, ind; + t.dval, t.val, (U, S, V⁺), dUSVᴴtrunc, ind; gauge_atol = gtol(dUSVᴴtrunc), degeneracy_atol = alg.val.rrule_alg.degeneracy_atol, ) end @@ -106,8 +109,9 @@ function EnzymeRules.augmented_primal( 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, alg_gauge, alg_fixed)) + return EnzymeRules.AugmentedReturn(primal, shadow, (env_conv, env, denv, alg_gauge, alg_fixed)) end function EnzymeRules.reverse( @@ -120,7 +124,7 @@ function EnzymeRules.reverse( 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.val, alg_gauge) + 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) @@ -150,6 +154,7 @@ function EnzymeRules.augmented_primal( end y = map(first, el_rrules) else + el_rrules = nothing y = map(a -> a.val, A) end dy = map(Enzyme.make_zero, y)