From 57250ffcb0eabd1ef778aa0de33f8ac6028347ae Mon Sep 17 00:00:00 2001 From: Yue Zhengyuan Date: Wed, 4 Nov 2026 21:03:47 +0800 Subject: [PATCH 01/11] Reuse VJP in GeomSum, ManualIter fixed-point gradient --- .../optimization/implicit_differentiation.jl | 35 ++++++++++--------- 1 file changed, 18 insertions(+), 17 deletions(-) diff --git a/src/algorithms/optimization/implicit_differentiation.jl b/src/algorithms/optimization/implicit_differentiation.jl index 37ae322b4..221212a7a 100644 --- a/src/algorithms/optimization/implicit_differentiation.jl +++ b/src/algorithms/optimization/implicit_differentiation.jl @@ -244,7 +244,11 @@ function _rrule( Δenv = unthunk(Δenv′) # evaluate the geometric sum - ∂F∂env = fixedpoint_gradient(Δenv, ∂f∂x, ∂f∂A, Δenv, gradmode.solver_alg) + ∂F∂env = if gradmode.solver_alg isa Union{GeomSum, ManualIter} + fixedpoint_gradient(Δenv, env_vjp, Δenv, gradmode.solver_alg) + else + fixedpoint_gradient(Δenv, ∂f∂x, ∂f∂A, Δenv, gradmode.solver_alg) + end return NoTangent(), ZeroTangent(), ∂F∂env, NoTangent() end @@ -254,10 +258,9 @@ end @doc raw""" fixedpoint_gradient(x̆, ∂ₓf, ∂ₚf, y₀, alg) + fixedpoint_gradient(x̆, vjp, y₀, alg::Union{GeomSum, ManualIter}) -Evaluates the VJP action ``x̆ ∂ₚx`` for an intermediate variable ``x \equiv x(p)`` -which satisfies the fixed-point equation ``x = f(x, p)``, given the VJP actions ``∂ₓf`` -and ``∂ₚf`` of the iterating function ``f``. +Evaluates the VJP action ``x̆ ∂ₚx`` for an intermediate variable ``x \equiv x(p)`` which satisfies the fixed-point equation ``x = f(x, p)``, given either the separate VJP actions ``∂ₓf`` and ``∂ₚf`` of the iterating function ``f`` or a shared pullback `vjp`. More specifically, given a cost function ``E(x(p), p)`` defined in terms of a set of variational parameters ``p`` and a set of intermediate variables ``x`` that depend on ``p``, @@ -278,18 +281,18 @@ This can be used to differentiate contraction routines, where ``p`` are the vari parameters of a tensor network, ``x̆ = ∂ₓE`` is the partial derivative of the cost function with respect to the contraction environment ``x``, and ``f`` is a single iteration of the contraction algorithm. + +`GeomSum` and `ManualIter` require the positional argument `vjp` in place of the separate VJP actions. +The shared pullback `vjp(y)` must return `(NoTangent(), ∂ₚf(y), ∂ₓf(y))`; the first entry is ignored. +The solver reuses both cotangents to avoid evaluating the same reverse pass twice. """ fixedpoint_gradient -# TODO: can we construct an implementation that does not need to evaluate the vjp -# twice if both ∂f∂A and ∂f∂x are needed? -function fixedpoint_gradient(∂E∂x, ∂f∂x, ∂f∂A, _, alg::GeomSum) - g = ∂E∂x - dx = ∂f∂A(g) # n = 0 term: ∂F∂x ∂f∂A +function fixedpoint_gradient(∂E∂x, vjp, _, alg::GeomSum) + _, dx, g = vjp(∂E∂x) # n = 0 term: ∂F∂x ∂f∂A ϵ = 2 * alg.tol for i in 1:(alg.maxiter) - g = ∂f∂x(g) - Σₙ = ∂f∂A(g) + _, Σₙ, g = vjp(g) dx += Σₙ ϵnew = norm(Σₙ) # TODO: normalize this error? Δϵ = ϵ - ϵnew @@ -305,20 +308,18 @@ function fixedpoint_gradient(∂E∂x, ∂f∂x, ∂f∂A, _, alg::GeomSum) return dx end -function fixedpoint_gradient(∂E∂x, ∂f∂x, ∂f∂A, y₀, alg::ManualIter) +function fixedpoint_gradient(∂E∂x, vjp, y₀, alg::ManualIter) y = deepcopy(y₀) # Do not mutate y₀ - dx = ∂f∂A(y) + _, dx, dy = vjp(y) ϵ = 1.0 for i in 1:(alg.maxiter) - y′ = ∂E∂x + ∂f∂x(y) - - dxnew = ∂f∂A(y′) + y′ = ∂E∂x + dy + _, dxnew, dy = vjp(y′) ϵnew = norm(dxnew - dx) Δϵ = ϵ - ϵnew alg.verbosity > 1 && @printf( "Gradient iter: %3d ‖Cᵢ₊₁-Cᵢ‖/N: %.2e Δ‖Cᵢ₊₁-Cᵢ‖/N: %.2e\n", i, ϵnew, Δϵ ) - y = y′ dx = dxnew ϵ = ϵnew From d8d845947733f44728d5b8634cadc11729e05c28 Mon Sep 17 00:00:00 2001 From: Yue Zhengyuan Date: Thu, 5 Nov 2026 16:26:49 +0800 Subject: [PATCH 02/11] =?UTF-8?q?Rename=20=E2=88=82f=E2=88=82A=20to=20?= =?UTF-8?q?=E2=88=82f=E2=88=82p?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../optimization/implicit_differentiation.jl | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/src/algorithms/optimization/implicit_differentiation.jl b/src/algorithms/optimization/implicit_differentiation.jl index 221212a7a..a213f742a 100644 --- a/src/algorithms/optimization/implicit_differentiation.jl +++ b/src/algorithms/optimization/implicit_differentiation.jl @@ -237,7 +237,7 @@ function _rrule( # prepare its pullback _, env_vjp = rrule_via_ad(config, gauge_fixed_iteration, state, env) # split off state and environment parts - ∂f∂A(x)::typeof(state) = env_vjp(x)[2] + ∂f∂p(x)::typeof(state) = env_vjp(x)[2] ∂f∂x(x)::typeof(env) = env_vjp(x)[3] function leading_boundary_fixed_pullback((Δenv′, Δinfo)) @@ -247,7 +247,7 @@ function _rrule( ∂F∂env = if gradmode.solver_alg isa Union{GeomSum, ManualIter} fixedpoint_gradient(Δenv, env_vjp, Δenv, gradmode.solver_alg) else - fixedpoint_gradient(Δenv, ∂f∂x, ∂f∂A, Δenv, gradmode.solver_alg) + fixedpoint_gradient(Δenv, ∂f∂x, ∂f∂p, Δenv, gradmode.solver_alg) end return NoTangent(), ZeroTangent(), ∂F∂env, NoTangent() @@ -289,7 +289,7 @@ The solver reuses both cotangents to avoid evaluating the same reverse pass twic fixedpoint_gradient function fixedpoint_gradient(∂E∂x, vjp, _, alg::GeomSum) - _, dx, g = vjp(∂E∂x) # n = 0 term: ∂F∂x ∂f∂A + _, dx, g = vjp(∂E∂x) # n = 0 term: ∂F∂x ∂f∂p ϵ = 2 * alg.tol for i in 1:(alg.maxiter) _, Σₙ, g = vjp(g) @@ -331,16 +331,16 @@ function fixedpoint_gradient(∂E∂x, vjp, y₀, alg::ManualIter) return dx end -function fixedpoint_gradient(∂E∂x, ∂f∂x, ∂f∂A, y₀, alg::KrylovKit.LinearSolver) +function fixedpoint_gradient(∂E∂x, ∂f∂x, ∂f∂p, y₀, alg::KrylovKit.LinearSolver) y, info = reallinsolve(∂f∂x, ∂E∂x, y₀, alg, 1, -1) if alg.verbosity > 0 && info.converged != 1 @warn("gradient fixed-point iteration reached maximal number of iterations:", info) end - return ∂f∂A(y) + return ∂f∂p(y) end -function fixedpoint_gradient(∂E∂x, ∂f∂x, ∂f∂A, x₀, alg::KrylovKit.KrylovAlgorithm) +function fixedpoint_gradient(∂E∂x, ∂f∂x, ∂f∂p, x₀, alg::KrylovKit.KrylovAlgorithm) function f(X) y = ∂f∂x(X[1]) return (VI.add!!(y, ∂E∂x, X[2]), X[2]) @@ -354,12 +354,12 @@ function fixedpoint_gradient(∂E∂x, ∂f∂x, ∂f∂A, x₀, alg::KrylovKit. @warn "Fixed-point gradient computation using Arnoldi failed:\n\tauxiliary component should be finite but was $(vecs[1][2])\n\tpossibly the Jacobian does not have a unique eigenvalue 1" @info "Falling back to linear solver for fixed-point gradient computation." backup_ls_alg = GMRES(; tol = alg.tol, maxiter = alg.maxiter, verbosity = alg.verbosity) - return fixedpoint_gradient(∂E∂x, ∂f∂x, ∂f∂A, x₀, backup_ls_alg) + return fixedpoint_gradient(∂E∂x, ∂f∂x, ∂f∂p, x₀, backup_ls_alg) else y = VI.scale!!(vecs[1][1], inv(vecs[1][2])) end - return ∂f∂A(y) + return ∂f∂p(y) end # From bd3888e702868f53ca143f796a5322d52cceca05 Mon Sep 17 00:00:00 2001 From: Yue Zhengyuan Date: Thu, 5 Nov 2026 16:31:04 +0800 Subject: [PATCH 03/11] =?UTF-8?q?Rename=20=CE=A3=E2=82=99=20to=20=CE=A3?= =?UTF-8?q?=E1=B5=A2?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../optimization/implicit_differentiation.jl | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/algorithms/optimization/implicit_differentiation.jl b/src/algorithms/optimization/implicit_differentiation.jl index a213f742a..6c8a20508 100644 --- a/src/algorithms/optimization/implicit_differentiation.jl +++ b/src/algorithms/optimization/implicit_differentiation.jl @@ -292,17 +292,17 @@ function fixedpoint_gradient(∂E∂x, vjp, _, alg::GeomSum) _, dx, g = vjp(∂E∂x) # n = 0 term: ∂F∂x ∂f∂p ϵ = 2 * alg.tol for i in 1:(alg.maxiter) - _, Σₙ, g = vjp(g) - dx += Σₙ - ϵnew = norm(Σₙ) # TODO: normalize this error? + _, Σᵢ, g = vjp(g) + dx += Σᵢ + ϵnew = norm(Σᵢ) # TODO: normalize this error? Δϵ = ϵ - ϵnew alg.verbosity > 1 && - @printf("Gradient iter: %3d ‖Σₙ‖: %.2e Δ‖Σₙ‖: %.2e\n", i, ϵnew, Δϵ) + @printf("Gradient iter: %3d ‖Σᵢ‖: %.2e Δ‖Σᵢ‖: %.2e\n", i, ϵnew, Δϵ) ϵ = ϵnew ϵ < alg.tol && break if alg.verbosity > 0 && i == alg.maxiter - @warn "gradient fixed-point iteration reached maximal number of iterations at ‖Σₙ‖ = $ϵ" + @warn "gradient fixed-point iteration reached maximal number of iterations at ‖Σᵢ‖ = $ϵ" end end return dx From a6c5ba62cfb18c4105799e98788ef43197b04c51 Mon Sep 17 00:00:00 2001 From: Yue Zhengyuan Date: Thu, 5 Nov 2026 16:58:00 +0800 Subject: [PATCH 04/11] Rename n to i --- src/algorithms/optimization/implicit_differentiation.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/algorithms/optimization/implicit_differentiation.jl b/src/algorithms/optimization/implicit_differentiation.jl index 6c8a20508..bbc17bcec 100644 --- a/src/algorithms/optimization/implicit_differentiation.jl +++ b/src/algorithms/optimization/implicit_differentiation.jl @@ -289,7 +289,7 @@ The solver reuses both cotangents to avoid evaluating the same reverse pass twic fixedpoint_gradient function fixedpoint_gradient(∂E∂x, vjp, _, alg::GeomSum) - _, dx, g = vjp(∂E∂x) # n = 0 term: ∂F∂x ∂f∂p + _, dx, g = vjp(∂E∂x) # i = 0 term: ∂E∂x ∂f∂p ϵ = 2 * alg.tol for i in 1:(alg.maxiter) _, Σᵢ, g = vjp(g) From ca9a7f9b98acad910cc6d5fc18f12c55f1436e01 Mon Sep 17 00:00:00 2001 From: Yue Zhengyuan Date: Thu, 5 Nov 2026 17:32:09 +0800 Subject: [PATCH 05/11] Unify `fixedpoint_gradient` signature --- .../optimization/implicit_differentiation.jl | 34 +++++++------------ 1 file changed, 13 insertions(+), 21 deletions(-) diff --git a/src/algorithms/optimization/implicit_differentiation.jl b/src/algorithms/optimization/implicit_differentiation.jl index bbc17bcec..7c8c77843 100644 --- a/src/algorithms/optimization/implicit_differentiation.jl +++ b/src/algorithms/optimization/implicit_differentiation.jl @@ -236,19 +236,12 @@ function _rrule( end # prepare its pullback _, env_vjp = rrule_via_ad(config, gauge_fixed_iteration, state, env) - # split off state and environment parts - ∂f∂p(x)::typeof(state) = env_vjp(x)[2] - ∂f∂x(x)::typeof(env) = env_vjp(x)[3] function leading_boundary_fixed_pullback((Δenv′, Δinfo)) Δenv = unthunk(Δenv′) - # evaluate the geometric sum - ∂F∂env = if gradmode.solver_alg isa Union{GeomSum, ManualIter} - fixedpoint_gradient(Δenv, env_vjp, Δenv, gradmode.solver_alg) - else - fixedpoint_gradient(Δenv, ∂f∂x, ∂f∂p, Δenv, gradmode.solver_alg) - end + # evaluate the fixed-point gradient + ∂F∂env = fixedpoint_gradient(Δenv, env_vjp, Δenv, gradmode.solver_alg) return NoTangent(), ZeroTangent(), ∂F∂env, NoTangent() end @@ -257,10 +250,10 @@ function _rrule( end @doc raw""" - fixedpoint_gradient(x̆, ∂ₓf, ∂ₚf, y₀, alg) - fixedpoint_gradient(x̆, vjp, y₀, alg::Union{GeomSum, ManualIter}) + fixedpoint_gradient(x̆, vjp, y₀, alg) -Evaluates the VJP action ``x̆ ∂ₚx`` for an intermediate variable ``x \equiv x(p)`` which satisfies the fixed-point equation ``x = f(x, p)``, given either the separate VJP actions ``∂ₓf`` and ``∂ₚf`` of the iterating function ``f`` or a shared pullback `vjp`. +Evaluates the VJP action ``x̆ ∂ₚx`` for an intermediate variable ``x \equiv x(p)`` which satisfies the fixed-point equation ``x = f(x, p)``. +The shared pullback `vjp(y)` of the iterating function ``f`` returns `(NoTangent(), ∂ₚf(y), ∂ₓf(y))`. More specifically, given a cost function ``E(x(p), p)`` defined in terms of a set of variational parameters ``p`` and a set of intermediate variables ``x`` that depend on ``p``, @@ -282,9 +275,7 @@ parameters of a tensor network, ``x̆ = ∂ₓE`` is the partial derivative of the cost function with respect to the contraction environment ``x``, and ``f`` is a single iteration of the contraction algorithm. -`GeomSum` and `ManualIter` require the positional argument `vjp` in place of the separate VJP actions. -The shared pullback `vjp(y)` must return `(NoTangent(), ∂ₚf(y), ∂ₓf(y))`; the first entry is ignored. -The solver reuses both cotangents to avoid evaluating the same reverse pass twice. +`GeomSum` and `ManualIter` reuse both cotangents at each step; the Krylov solvers use the environment cotangent during the solve and evaluate the parameter cotangent once afterward. """ fixedpoint_gradient @@ -331,18 +322,19 @@ function fixedpoint_gradient(∂E∂x, vjp, y₀, alg::ManualIter) return dx end -function fixedpoint_gradient(∂E∂x, ∂f∂x, ∂f∂p, y₀, alg::KrylovKit.LinearSolver) +function fixedpoint_gradient(∂E∂x, vjp, y₀, alg::KrylovKit.LinearSolver) + ∂f∂x(x) = vjp(x)[3] y, info = reallinsolve(∂f∂x, ∂E∂x, y₀, alg, 1, -1) if alg.verbosity > 0 && info.converged != 1 @warn("gradient fixed-point iteration reached maximal number of iterations:", info) end - return ∂f∂p(y) + return vjp(y)[2] end -function fixedpoint_gradient(∂E∂x, ∂f∂x, ∂f∂p, x₀, alg::KrylovKit.KrylovAlgorithm) +function fixedpoint_gradient(∂E∂x, vjp, x₀, alg::KrylovKit.KrylovAlgorithm) function f(X) - y = ∂f∂x(X[1]) + y = vjp(X[1])[3] return (VI.add!!(y, ∂E∂x, X[2]), X[2]) end X₀ = (x₀, one(scalartype(x₀))) @@ -354,12 +346,12 @@ function fixedpoint_gradient(∂E∂x, ∂f∂x, ∂f∂p, x₀, alg::KrylovKit. @warn "Fixed-point gradient computation using Arnoldi failed:\n\tauxiliary component should be finite but was $(vecs[1][2])\n\tpossibly the Jacobian does not have a unique eigenvalue 1" @info "Falling back to linear solver for fixed-point gradient computation." backup_ls_alg = GMRES(; tol = alg.tol, maxiter = alg.maxiter, verbosity = alg.verbosity) - return fixedpoint_gradient(∂E∂x, ∂f∂x, ∂f∂p, x₀, backup_ls_alg) + return fixedpoint_gradient(∂E∂x, vjp, x₀, backup_ls_alg) else y = VI.scale!!(vecs[1][1], inv(vecs[1][2])) end - return ∂f∂p(y) + return vjp(y)[2] end # From 283315c1efb50c538f79218111c5df978221ff98 Mon Sep 17 00:00:00 2001 From: Yue Zhengyuan Date: Fri, 6 Nov 2026 10:29:52 +0800 Subject: [PATCH 06/11] Apply review suggestions --- .../optimization/implicit_differentiation.jl | 91 ++++++++++--------- 1 file changed, 46 insertions(+), 45 deletions(-) diff --git a/src/algorithms/optimization/implicit_differentiation.jl b/src/algorithms/optimization/implicit_differentiation.jl index 7c8c77843..445945a53 100644 --- a/src/algorithms/optimization/implicit_differentiation.jl +++ b/src/algorithms/optimization/implicit_differentiation.jl @@ -227,35 +227,38 @@ function _rrule( alg_gauge = _scrambling_env_gauge(alg) # select appropriate gauge-fixing algorithm env_conv, info = ctmrg_iteration(InfiniteSquareNetwork(state), env, alg_fixed) signs, = compute_gauge_fix_gauge(env_conv, env, alg_gauge) - # fix global phases relative to the fixed point inside the iteration - function gauge_fixed_iteration(A, x) + # `p` for state, `x` for environment + function gauge_fixed_iteration(p, x) + # fix global phases relative to the fixed point inside the iteration return fix_global_phases( - fix_relative_phases(ctmrg_iteration(InfiniteSquareNetwork(A), x, alg_fixed)[1], signs), + fix_relative_phases(ctmrg_iteration(InfiniteSquareNetwork(p), x, alg_fixed)[1], signs), env, ) end # prepare its pullback - _, env_vjp = rrule_via_ad(config, gauge_fixed_iteration, state, env) + _, ∂f = rrule_via_ad(config, gauge_fixed_iteration, state, env) + + # TODO: try partial-argument VJPs after Enzyme support is added. + # fₚ = Base.Fix1(gauge_fixed_iteration, state) # fix state `p` + # fₓ = Base.Fix2(gauge_fixed_iteration, env) # fix environment `x` function leading_boundary_fixed_pullback((Δenv′, Δinfo)) Δenv = unthunk(Δenv′) - # evaluate the fixed-point gradient - ∂F∂env = fixedpoint_gradient(Δenv, env_vjp, Δenv, gradmode.solver_alg) - - return NoTangent(), ZeroTangent(), ∂F∂env, NoTangent() + ∂E∂env = fixedpoint_gradient(Δenv, ∂f, Δenv, gradmode.solver_alg) + return NoTangent(), ZeroTangent(), ∂E∂env, NoTangent() end return (env, info), leading_boundary_fixed_pullback end @doc raw""" - fixedpoint_gradient(x̆, vjp, y₀, alg) + fixedpoint_gradient(x̆, ∂f, y₀, alg) -Evaluates the VJP action ``x̆ ∂ₚx`` for an intermediate variable ``x \equiv x(p)`` which satisfies the fixed-point equation ``x = f(x, p)``. -The shared pullback `vjp(y)` of the iterating function ``f`` returns `(NoTangent(), ∂ₚf(y), ∂ₓf(y))`. +Evaluates the VJP action ``x̆ ∂ₚx`` for an intermediate variable ``x \equiv x(p)`` which satisfies the fixed-point equation ``x = f(p, x)``. +The shared pullback `∂f(y)` of the iterating function ``f`` returns `(NoTangent(), ∂ₚf(y), ∂ₓf(y))`. -More specifically, given a cost function ``E(x(p), p)`` defined in terms of a set of +More specifically, given a cost function ``E(p, x(p))`` defined in terms of a set of variational parameters ``p`` and a set of intermediate variables ``x`` that depend on ``p``, ``x \equiv x(p)``, the gradient of the cost function is given by @@ -263,9 +266,7 @@ variational parameters ``p`` and a set of intermediate variables ``x`` that depe dE/dp = ∂ₓE ∂ₚx + ∂ₚE. ``` -Given the fixed-point equation ``x = f(x, p)``, the VJP action of the Jacobian ``∂ₚx``` on -the adjoint ``x̆ = ∂ₓE`` in the first term of this expression can be evaluated through -implicit differentiation of the fixed-point condition as +Given the fixed-point equation ``x = f(p, x)``, the VJP action of the Jacobian ``∂ₚx`` on the adjoint ``x̆ = ∂ₓE`` in the first term of this expression can be evaluated through implicit differentiation of the fixed-point condition as ```math x̆ ∂ₚx = x̆ (1 - ∂ₓf)⁻¹ ∂ₚf = ∑ₙ x̆ (∂ₓf)ⁿ ∂ₚf. ``` @@ -274,16 +275,14 @@ This can be used to differentiate contraction routines, where ``p`` are the vari parameters of a tensor network, ``x̆ = ∂ₓE`` is the partial derivative of the cost function with respect to the contraction environment ``x``, and ``f`` is a single iteration of the contraction algorithm. - -`GeomSum` and `ManualIter` reuse both cotangents at each step; the Krylov solvers use the environment cotangent during the solve and evaluate the parameter cotangent once afterward. """ fixedpoint_gradient -function fixedpoint_gradient(∂E∂x, vjp, _, alg::GeomSum) - _, dx, g = vjp(∂E∂x) # i = 0 term: ∂E∂x ∂f∂p +function fixedpoint_gradient(∂E∂x, ∂f, _, alg::GeomSum) + _, dx, g = ∂f(∂E∂x) # i = 0 term: ∂E∂x ∂f∂p ϵ = 2 * alg.tol for i in 1:(alg.maxiter) - _, Σᵢ, g = vjp(g) + _, Σᵢ, g = ∂f(g) dx += Σᵢ ϵnew = norm(Σᵢ) # TODO: normalize this error? Δϵ = ϵ - ϵnew @@ -299,13 +298,13 @@ function fixedpoint_gradient(∂E∂x, vjp, _, alg::GeomSum) return dx end -function fixedpoint_gradient(∂E∂x, vjp, y₀, alg::ManualIter) +function fixedpoint_gradient(∂E∂x, ∂f, y₀, alg::ManualIter) y = deepcopy(y₀) # Do not mutate y₀ - _, dx, dy = vjp(y) + _, dx, dy = ∂f(y) ϵ = 1.0 for i in 1:(alg.maxiter) y′ = ∂E∂x + dy - _, dxnew, dy = vjp(y′) + _, dxnew, dy = ∂f(y′) ϵnew = norm(dxnew - dx) Δϵ = ϵ - ϵnew alg.verbosity > 1 && @printf( @@ -322,36 +321,38 @@ function fixedpoint_gradient(∂E∂x, vjp, y₀, alg::ManualIter) return dx end -function fixedpoint_gradient(∂E∂x, vjp, y₀, alg::KrylovKit.LinearSolver) - ∂f∂x(x) = vjp(x)[3] +function fixedpoint_gradient(∂E∂x, ∂f, y₀, alg::KrylovKit.LinearSolver) + ∂f∂p(x) = ∂f(x)[2] + ∂f∂x(x) = ∂f(x)[3] y, info = reallinsolve(∂f∂x, ∂E∂x, y₀, alg, 1, -1) if alg.verbosity > 0 && info.converged != 1 @warn("gradient fixed-point iteration reached maximal number of iterations:", info) end - - return vjp(y)[2] + return ∂f∂p(y) end -function fixedpoint_gradient(∂E∂x, vjp, x₀, alg::KrylovKit.KrylovAlgorithm) - function f(X) - y = vjp(X[1])[3] - return (VI.add!!(y, ∂E∂x, X[2]), X[2]) +function fixedpoint_gradient(∂E∂x, ∂f, y₀, alg::KrylovKit.KrylovAlgorithm) + ∂f∂p(x) = ∂f(x)[2] + ∂f∂x(x) = ∂f(x)[3] + function augmented_map(Y) + local y, s = Y + return (VI.add!!(∂f∂x(y), ∂E∂x, s), s) end - X₀ = (x₀, one(scalartype(x₀))) - _, vecs, info = realeigsolve(f, X₀, 1, :LM, alg) + Y₀ = (y₀, one(scalartype(y₀))) + _, Ys, info = realeigsolve(augmented_map, Y₀, 1, :LM, alg) if alg.verbosity > 0 && info.converged < 1 @warn("gradient fixed-point iteration reached maximal number of iterations:", info) end - if norm(vecs[1][2]) < 1.0e-2 * alg.tol - @warn "Fixed-point gradient computation using Arnoldi failed:\n\tauxiliary component should be finite but was $(vecs[1][2])\n\tpossibly the Jacobian does not have a unique eigenvalue 1" + y, s = first(Ys) + if norm(s) < 1.0e-2 * alg.tol + @warn "Fixed-point gradient computation using Arnoldi failed:\n\tauxiliary component should be finite but was $(s)\n\tpossibly the Jacobian does not have a unique eigenvalue 1" @info "Falling back to linear solver for fixed-point gradient computation." backup_ls_alg = GMRES(; tol = alg.tol, maxiter = alg.maxiter, verbosity = alg.verbosity) - return fixedpoint_gradient(∂E∂x, vjp, x₀, backup_ls_alg) + return fixedpoint_gradient(∂E∂x, ∂f, y₀, backup_ls_alg) else - y = VI.scale!!(vecs[1][1], inv(vecs[1][2])) + y = VI.scale!!(y, inv(s)) end - - return vjp(y)[2] + return ∂f∂p(y) end # @@ -423,11 +424,11 @@ end @doc raw""" implicit_gradient(x̆, ∂ₓF, ∂ₚF, y₀, alg) -Evaluates the VJP action ``x̆ ∂ₚx`` for an intermediate variable ``x \equiv x(p)`` -which satisfies the characteristic equation ``F(x, p) = 0``, given the VJP actions ``∂ₓF`` +Evaluates the VJP action ``x̆ ∂ₚx`` for an intermediate variable ``x \equiv x(p)`` +which satisfies the characteristic equation ``F(p, x) = 0``, given the VJP actions ``∂ₓF`` and ``∂ₚF`` of the function ``F`` encoding the characteristic equation. -More specifically, given a cost function ``E(x(p), p)`` defined in terms of a set of +More specifically, given a cost function ``E(p, x(p))`` defined in terms of a set of variational parameters ``p`` and a set of intermediate variables ``x`` that depend on ``p``, ``x \equiv x(p)``, the gradient of the cost function is given by @@ -435,7 +436,7 @@ variational parameters ``p`` and a set of intermediate variables ``x`` that depe dE/dp = ∂ₓE ∂ₚx + ∂ₚE. ``` -Given the characteristic equation ``F(x, p) = 0``, the VJP action of the Jacobian ``∂ₚx``` on +Given the characteristic equation ``F(p, x) = 0``, the VJP action of the Jacobian ``∂ₚx``` on the adjoint ``x̆ = ∂ₓE`` in the first term of this expression can be evaluated through implicit differentiation of the characteristic equation as ```math @@ -449,13 +450,13 @@ is an algebraic function characterizing the convergence of the algorithm. """ implicit_gradient -function implicit_gradient(∂E∂x, ∂F∂x, ∂F∂A, y₀, alg::KrylovKit.LinearSolver) +function implicit_gradient(∂E∂x, ∂F∂x, ∂F∂p, y₀, alg::KrylovKit.LinearSolver) y, info = reallinsolve(∂F∂x, ∂E∂x, y₀, alg) if alg.verbosity > 0 && info.converged != 1 @warn("Implicit gradient linear solve reached maximal number of iterations:", info) end - return (-1) * ∂F∂A(y) + return (-1) * ∂F∂p(y) end ## C4vCTMRG gradient through implicit differentiation From d3e9ee47fa8900e957bd654200f10a84560fcb29 Mon Sep 17 00:00:00 2001 From: Yue Zhengyuan Date: Fri, 6 Nov 2026 11:05:20 +0800 Subject: [PATCH 07/11] Docstring rewrite --- .../optimization/implicit_differentiation.jl | 24 +++++++++---------- 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/src/algorithms/optimization/implicit_differentiation.jl b/src/algorithms/optimization/implicit_differentiation.jl index 445945a53..98ee37e84 100644 --- a/src/algorithms/optimization/implicit_differentiation.jl +++ b/src/algorithms/optimization/implicit_differentiation.jl @@ -255,26 +255,26 @@ end @doc raw""" fixedpoint_gradient(x̆, ∂f, y₀, alg) -Evaluates the VJP action ``x̆ ∂ₚx`` for an intermediate variable ``x \equiv x(p)`` which satisfies the fixed-point equation ``x = f(p, x)``. -The shared pullback `∂f(y)` of the iterating function ``f`` returns `(NoTangent(), ∂ₚf(y), ∂ₓf(y))`. +Given the iteration function ``f(p, x)``, with ``p`` as parameters and ``x \equiv x(p)`` as the solution of the fixed-point equation ``x = f(p, x)``, evaluates the VJP action ``x̆ ∂ₚx``, where ``x̆`` is a covector. -More specifically, given a cost function ``E(p, x(p))`` defined in terms of a set of -variational parameters ``p`` and a set of intermediate variables ``x`` that depend on ``p``, -``x \equiv x(p)``, the gradient of the cost function is given by +In the context of cost function ``E(p, x(p))`` optimization, let ``x̆ = ∂ₓE``, ``p̆ = ∂ₚE``. +The gradient of the cost function is given by ```math -dE/dp = ∂ₓE ∂ₚx + ∂ₚE. +dE/dp = x̆ ∂ₚx + p̆. ``` -Given the fixed-point equation ``x = f(p, x)``, the VJP action of the Jacobian ``∂ₚx`` on the adjoint ``x̆ = ∂ₓE`` in the first term of this expression can be evaluated through implicit differentiation of the fixed-point condition as +Implicit differentiation of the fixed-point equation gives + ```math -x̆ ∂ₚx = x̆ (1 - ∂ₓf)⁻¹ ∂ₚf = ∑ₙ x̆ (∂ₓf)ⁿ ∂ₚf. +x̆ ∂ₚx = y ∂ₚf, \quad y = x̆ (1 - ∂ₓf)⁻¹. ``` -This can be used to differentiate contraction routines, where ``p`` are the variational -parameters of a tensor network, ``x̆ = ∂ₓE`` is the partial -derivative of the cost function with respect to the contraction environment ``x``, and ``f`` -is a single iteration of the contraction algorithm. +- The pullback `∂f(y)` returns `(NoTangent(), ∂ₚf(y), ∂ₓf(y))`. +- `alg` is the algorithm to solve ``y`` from the equation ``y (1 - ∂ₓf) = x̆``, with initial guess `y₀` (unused by `GeomSum` algorithm). +- The contribution ``p̆ = ∂ₚE`` should be added separately. + +For tensor-network contraction, ``p`` are the variational parameters, ``x`` is the contraction environment, and ``f`` is one iteration of the contraction algorithm. """ fixedpoint_gradient From 4165ffa5fed1ac98eb4306a8b9e2b2d9f85c5b86 Mon Sep 17 00:00:00 2001 From: Yue Zhengyuan Date: Fri, 6 Nov 2026 11:19:14 +0800 Subject: [PATCH 08/11] Add geometric summation back in docstring --- src/algorithms/optimization/implicit_differentiation.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/algorithms/optimization/implicit_differentiation.jl b/src/algorithms/optimization/implicit_differentiation.jl index 98ee37e84..5373d0373 100644 --- a/src/algorithms/optimization/implicit_differentiation.jl +++ b/src/algorithms/optimization/implicit_differentiation.jl @@ -267,12 +267,12 @@ dE/dp = x̆ ∂ₚx + p̆. Implicit differentiation of the fixed-point equation gives ```math -x̆ ∂ₚx = y ∂ₚf, \quad y = x̆ (1 - ∂ₓf)⁻¹. +x̆ ∂ₚx = y ∂ₚf, \quad y = x̆ (1 - ∂ₓf)⁻¹ = ∑ₙ x̆ (∂ₓf)ⁿ. ``` - The pullback `∂f(y)` returns `(NoTangent(), ∂ₚf(y), ∂ₓf(y))`. - `alg` is the algorithm to solve ``y`` from the equation ``y (1 - ∂ₓf) = x̆``, with initial guess `y₀` (unused by `GeomSum` algorithm). -- The contribution ``p̆ = ∂ₚE`` should be added separately. +- The contribution ``p̆ = ∂ₚE`` in ``dE/dp`` should be added separately. For tensor-network contraction, ``p`` are the variational parameters, ``x`` is the contraction environment, and ``f`` is one iteration of the contraction algorithm. """ From 4ed43fdf5e784c6352893ec757e79ee094898062 Mon Sep 17 00:00:00 2001 From: Yue Zhengyuan Date: Fri, 6 Nov 2026 12:07:37 +0800 Subject: [PATCH 09/11] =?UTF-8?q?Type=20annotate=20=E2=88=82f?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/algorithms/optimization/implicit_differentiation.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/algorithms/optimization/implicit_differentiation.jl b/src/algorithms/optimization/implicit_differentiation.jl index 5373d0373..ad36c3310 100644 --- a/src/algorithms/optimization/implicit_differentiation.jl +++ b/src/algorithms/optimization/implicit_differentiation.jl @@ -236,7 +236,8 @@ function _rrule( ) end # prepare its pullback - _, ∂f = rrule_via_ad(config, gauge_fixed_iteration, state, env) + _, pullback = rrule_via_ad(config, gauge_fixed_iteration, state, env) + ∂f = x -> pullback(x)::Tuple{Any, typeof(state), typeof(env)} # TODO: try partial-argument VJPs after Enzyme support is added. # fₚ = Base.Fix1(gauge_fixed_iteration, state) # fix state `p` From 9f2df64f3b6bcf4679b67b4df70d41cb4a016258 Mon Sep 17 00:00:00 2001 From: Yue Zhengyuan Date: Fri, 6 Nov 2026 20:10:41 +0800 Subject: [PATCH 10/11] =?UTF-8?q?Better=20way=20to=20annotate=20type=20of?= =?UTF-8?q?=20=E2=88=82f?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/algorithms/optimization/implicit_differentiation.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/algorithms/optimization/implicit_differentiation.jl b/src/algorithms/optimization/implicit_differentiation.jl index ad36c3310..b6d9c540d 100644 --- a/src/algorithms/optimization/implicit_differentiation.jl +++ b/src/algorithms/optimization/implicit_differentiation.jl @@ -237,7 +237,7 @@ function _rrule( end # prepare its pullback _, pullback = rrule_via_ad(config, gauge_fixed_iteration, state, env) - ∂f = x -> pullback(x)::Tuple{Any, typeof(state), typeof(env)} + ∂f(y)::Tuple{Any, typeof(state), typeof(env)} = pullback(y) # TODO: try partial-argument VJPs after Enzyme support is added. # fₚ = Base.Fix1(gauge_fixed_iteration, state) # fix state `p` From 09e9b954925fb70fd8e5f27541b8b371d9666e3f Mon Sep 17 00:00:00 2001 From: Yue Zhengyuan Date: Fri, 6 Nov 2026 20:11:39 +0800 Subject: [PATCH 11/11] =?UTF-8?q?Correct=20description=20of=20=E2=88=82f?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/algorithms/optimization/implicit_differentiation.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/algorithms/optimization/implicit_differentiation.jl b/src/algorithms/optimization/implicit_differentiation.jl index b6d9c540d..111699453 100644 --- a/src/algorithms/optimization/implicit_differentiation.jl +++ b/src/algorithms/optimization/implicit_differentiation.jl @@ -271,7 +271,7 @@ Implicit differentiation of the fixed-point equation gives x̆ ∂ₚx = y ∂ₚf, \quad y = x̆ (1 - ∂ₓf)⁻¹ = ∑ₙ x̆ (∂ₓf)ⁿ. ``` -- The pullback `∂f(y)` returns `(NoTangent(), ∂ₚf(y), ∂ₓf(y))`. +- The second and third components of the pullback `∂f(y)` are the VJP actions ``y ∂ₚf`` and ``y ∂ₓf``, respectively. The first component is ignored. - `alg` is the algorithm to solve ``y`` from the equation ``y (1 - ∂ₓf) = x̆``, with initial guess `y₀` (unused by `GeomSum` algorithm). - The contribution ``p̆ = ∂ₚE`` in ``dE/dp`` should be added separately.