From a2c55dada3e05d53028b3cc883fdf23e6cb78eaa Mon Sep 17 00:00:00 2001 From: Yue Zhengyuan Date: Sat, 17 Oct 2026 09:29:32 +0800 Subject: [PATCH 1/4] Add `LQAdjoint` for `right_orth` --- src/PEPSKit.jl | 3 ++- src/utility/lq.jl | 39 ++++++++++++++++++++++++++++ src/utility/qr.jl | 10 ++++--- test/testsuite/TestSuite.jl | 6 +++++ test/testsuite/utility/qr_wrapper.jl | 36 +++++++++++++++++++++++++ test/utility/qr_wrapper.jl | 11 ++++++++ 6 files changed, 101 insertions(+), 4 deletions(-) create mode 100644 src/utility/lq.jl create mode 100644 test/testsuite/utility/qr_wrapper.jl create mode 100644 test/utility/qr_wrapper.jl diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl index b57b6361f..ebc20f091 100644 --- a/src/PEPSKit.jl +++ b/src/PEPSKit.jl @@ -63,6 +63,7 @@ include("utility/twistdual.jl") include("utility/eigh.jl") include("utility/svd.jl") include("utility/qr.jl") +include("utility/lq.jl") include("utility/rotations.jl") include("utility/hook_pullback.jl") include("utility/autoopt.jl") @@ -172,7 +173,7 @@ include("algorithms/select_algorithm.jl") using .Defaults: set_scheduler! export set_scheduler! -export EighAdjoint, IterEigh, SVDAdjoint, IterSVD, QRAdjoint +export EighAdjoint, IterEigh, SVDAdjoint, IterSVD, QRAdjoint, LQAdjoint export CTMRGEnv, SequentialCTMRG, SimultaneousCTMRG export initialize_ctmrg_environment, RandomInitialization, ProductStateInitialization, ApplicationInitialization, IdentityInitialization diff --git a/src/utility/lq.jl b/src/utility/lq.jl new file mode 100644 index 000000000..37396c194 --- /dev/null +++ b/src/utility/lq.jl @@ -0,0 +1,39 @@ +""" +$(TYPEDEF) + +Wrapper for an LQ decomposition algorithm with a defined reverse rule, used by `right_orth(!)`. +See [`QRAdjoint`](@ref) for details of the constructor options and defaults. +""" +struct LQAdjoint{F, R} + fwd_alg::F + rrule_alg::R +end + +LQAdjoint(; kwargs...) = LQAdjoint(_orthogonalization_algorithms(; kwargs...)...) + +"""Compute an LQ decomposition using the forward algorithm and reverse rule configured by `LQAdjoint`.""" +MatrixAlgebraKit.right_orth(t, alg::LQAdjoint) = right_orth!(copy(t), alg) +"""Factorize the input in place using the LQ forward algorithm configured by `LQAdjoint`.""" +MatrixAlgebraKit.right_orth!(t, alg::LQAdjoint) = right_orth!(t, alg.fwd_alg) + +# right_orth! rrule wrapping MatrixAlgebraKit's lq_pullback! +function ChainRulesCore.rrule( + ::typeof(right_orth!), + t::AbstractTensorMap, + alg::LQAdjoint{F, R}, + ) where {F <: MatrixAlgebraKit.Algorithm, R <: FullPullback} + LQ = right_orth(t, alg) + gtol = _get_pullback_gauge_tol(alg.rrule_alg.verbosity) + + """Propagate LQ cotangents with the configured gauge-sensitivity tolerance.""" + function right_orth!_pullback(ΔLQ) + Δt = zeros(scalartype(t), space(t)) + MatrixAlgebraKit.lq_pullback!(Δt, t, LQ, unthunk.(ΔLQ); gauge_atol = gtol(ΔLQ)) + return NoTangent(), Δt, NoTangent() + end + function right_orth!_pullback(::Tuple{ZeroTangent, ZeroTangent}) + return NoTangent(), ZeroTangent(), NoTangent() + end + + return LQ, right_orth!_pullback +end diff --git a/src/utility/qr.jl b/src/utility/qr.jl index a512821c8..194a5f5f6 100644 --- a/src/utility/qr.jl +++ b/src/utility/qr.jl @@ -19,7 +19,8 @@ Construct a `QRAdjoint` algorithm struct based on the following keyword argument The available algorithms are provided through MatrixAlgebraKit and include: - `:DefaultAlgorithm` : MatrixAlgebraKit's [default QR algorithm](@extref MatrixAlgebraKit.DefaultAlgorithm) for a given matrix type. - `:Householder` : MatrixAlgebraKit's [`Householder`](@extref MatrixAlgebraKit.Householder) -* `rrule_alg::Union{Algorithm,NamedTuple}=(; alg::Symbol=$(Defaults.qr_rrule_alg))`: Reverse-rule algorithm for differentiating the eigenvalue decomposition. Can be supplied by an `Algorithm` instance directly or as a `NamedTuple` where `alg` is one of the following: +* `rrule_alg::Union{Algorithm,NamedTuple}=(; alg::Symbol=$(Defaults.qr_rrule_alg))`: Reverse-rule algorithm for differentiating the QR decomposition. + Can be supplied by an `Algorithm` instance directly or as a `NamedTuple` where `alg` is one of the following: - `:FullPullback` : MatrixAlgebraKit's [`qr_pullback!`](@extref MatrixAlgebraKit.qr_pullback!) !!! note @@ -39,7 +40,10 @@ const QR_RRULE_SYMBOLS = IdDict{Symbol, Type{<:Any}}( :FullPullback => FullPullback ) -function QRAdjoint(; fwd_alg = (;), rrule_alg = (;)) +QRAdjoint(; kwargs...) = QRAdjoint(_orthogonalization_algorithms(; kwargs...)...) + +"""Parse the shared forward and reverse algorithm options for QR and LQ decompositions.""" +function _orthogonalization_algorithms(; fwd_alg = (;), rrule_alg = (;)) # parse forward algorithm fwd_algorithm = if fwd_alg isa NamedTuple fwd_kwargs = (; @@ -72,7 +76,7 @@ function QRAdjoint(; fwd_alg = (;), rrule_alg = (;)) rrule_alg end - return QRAdjoint(fwd_algorithm, rrule_algorithm) + return fwd_algorithm, rrule_algorithm end """ diff --git a/test/testsuite/TestSuite.jl b/test/testsuite/TestSuite.jl index db2c7b000..fd4223c5b 100644 --- a/test/testsuite/TestSuite.jl +++ b/test/testsuite/TestSuite.jl @@ -266,6 +266,12 @@ module UtilityEighWrapper end using .UtilityEighWrapper +module UtilityQRWrapper + include("utility/qr_wrapper.jl") + export utility_qr_wrapper +end +using .UtilityQRWrapper + module UtilitySVDWrapper include("utility/svd_wrapper.jl") export utility_svd_wrapper diff --git a/test/testsuite/utility/qr_wrapper.jl b/test/testsuite/utility/qr_wrapper.jl new file mode 100644 index 000000000..9e818d2e4 --- /dev/null +++ b/test/testsuite/utility/qr_wrapper.jl @@ -0,0 +1,36 @@ +using Test, Random, LinearAlgebra +using TensorKit, MatrixAlgebraKit, PEPSKit, Zygote, Adapt +using ChainRulesCore: rrule, NoTangent, ZeroTangent + +"""Check LQ factors and pullbacks against adjointed QR, including mixed dual fermionic legs.""" +function utility_qr_wrapper(AT) + Random.seed!(73194) + qr_alg, lq_alg = QRAdjoint(), LQAdjoint() + factorize = X -> right_orth(X, lq_alg) + reference = X -> reverse(adjoint.(left_orth(X', qr_alg))) + spaces = (ComplexSpace(2), Vect[FermionParity](0 => 1, 1 => 1), Vect[FermionNumber](-1 => 1, 0 => 1, 1 => 1)) + return @testset "QR/LQ wrapper ($(sectortype(V)), $nout ← $nin, $AT)" for V in spaces, (nout, nin) in ((2, 3), (3, 2)) + output, input = (prod(ntuple(i -> isodd(i) ? V : V', n)) for n in (nout, nin)) + A = adapt(AT, randn(ComplexF64, output ← input)) + L, Q = factorize(A) + @test all(isapprox.((L, Q), reference(A))) + @test L * Q ≈ A + @test right_orth!(copy(A), lq_alg) == (L, Q) + + H, K = (adapt(AT, randn(ComplexF64, W ← W)) for W in (output, input)) + # The real Gram-matrix loss is gauge invariant and exercises both factors. + loss = decomposition -> X -> let (l, q) = decomposition(X) + real(dot(H, l * l') + dot(K, q' * q)) + end + f, fref = loss(factorize), loss(reference) + gradient = only(Zygote.gradient(f, A)) + @test gradient ≈ only(Zygote.gradient(fref, A)) + direction = normalize(adapt(AT, randn(ComplexF64, space(A)))) + epsilon = 1.0e-5 + finite = (f(A + epsilon * direction) - f(A - epsilon * direction)) / (2 * epsilon) + @test real(dot(gradient, direction)) ≈ finite atol = 1.0e-7 rtol = 1.0e-7 + + _, back = rrule(right_orth!, copy(A), lq_alg) + @test back((ZeroTangent(), ZeroTangent())) == (NoTangent(), ZeroTangent(), NoTangent()) + end +end diff --git a/test/utility/qr_wrapper.jl b/test/utility/qr_wrapper.jl new file mode 100644 index 000000000..f39bea571 --- /dev/null +++ b/test/utility/qr_wrapper.jl @@ -0,0 +1,11 @@ +using Test +using PEPSKit + +@isdefined(TestSuite) || include("../testsuite/TestSuite.jl") +using .TestSuite + +is_buildkite = get(ENV, "BUILDKITE", "false") == "true" + +if !is_buildkite + TestSuite.utility_qr_wrapper(Vector) +end From 7f148adf9905d6c83380f770feec68de0e46d1bd Mon Sep 17 00:00:00 2001 From: Yue Zhengyuan Date: Sat, 17 Oct 2026 13:37:56 +0800 Subject: [PATCH 2/4] Address test timeout --- test/testsuite/utility/qr_wrapper.jl | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/test/testsuite/utility/qr_wrapper.jl b/test/testsuite/utility/qr_wrapper.jl index 9e818d2e4..348cc66ce 100644 --- a/test/testsuite/utility/qr_wrapper.jl +++ b/test/testsuite/utility/qr_wrapper.jl @@ -7,7 +7,10 @@ function utility_qr_wrapper(AT) Random.seed!(73194) qr_alg, lq_alg = QRAdjoint(), LQAdjoint() factorize = X -> right_orth(X, lq_alg) - reference = X -> reverse(adjoint.(left_orth(X', qr_alg))) + # Explicit tuple adjoints avoid costly inference in the broadcast pullback. + reference = X -> let (q, r) = left_orth(X', qr_alg) + (r', q') + end spaces = (ComplexSpace(2), Vect[FermionParity](0 => 1, 1 => 1), Vect[FermionNumber](-1 => 1, 0 => 1, 1 => 1)) return @testset "QR/LQ wrapper ($(sectortype(V)), $nout ← $nin, $AT)" for V in spaces, (nout, nin) in ((2, 3), (3, 2)) output, input = (prod(ntuple(i -> isodd(i) ? V : V', n)) for n in (nout, nin)) From ddef98d62b7368681fa854c678c8b087a263b447 Mon Sep 17 00:00:00 2001 From: Yue Zhengyuan Date: Mon, 19 Oct 2026 21:01:59 +0800 Subject: [PATCH 3/4] Unified `OrthAdjoint` --- src/Defaults.jl | 10 +- src/PEPSKit.jl | 5 +- src/algorithms/ctmrg/c4v.jl | 2 +- src/algorithms/ctmrg/ctmrg.jl | 8 +- src/algorithms/ctmrg/projectors/c4v_qr.jl | 2 +- src/algorithms/ctmrg/projectors/projectors.jl | 2 +- src/utility/lq.jl | 39 ----- src/utility/orth.jl | 139 ++++++++++++++++++ src/utility/qr.jl | 112 -------------- test/testsuite/TestSuite.jl | 8 +- .../gradients/c4v_ctmrg_gradients.jl | 2 +- test/testsuite/utility/orth_wrapper.jl | 56 +++++++ test/testsuite/utility/qr_wrapper.jl | 39 ----- .../{qr_wrapper.jl => orth_wrapper.jl} | 2 +- 14 files changed, 214 insertions(+), 212 deletions(-) delete mode 100644 src/utility/lq.jl create mode 100644 src/utility/orth.jl delete mode 100644 src/utility/qr.jl create mode 100644 test/testsuite/utility/orth_wrapper.jl delete mode 100644 test/testsuite/utility/qr_wrapper.jl rename test/utility/{qr_wrapper.jl => orth_wrapper.jl} (81%) diff --git a/src/Defaults.jl b/src/Defaults.jl index da8afab9a..c1537907d 100644 --- a/src/Defaults.jl +++ b/src/Defaults.jl @@ -162,11 +162,11 @@ const eigh_fwd_alg = :DefaultAlgorithm # ∈ {: const eigh_rrule_alg = :FullPullback # ∈ {:FullPullback, :TruncPullback} const eigh_rrule_verbosity = 0 -# QR forward & reverse -const qr_fwd_alg = :DefaultAlgorithm -const qr_fwd_positive = true -const qr_rrule_alg = :FullPullback -const qr_rrule_verbosity = 0 +# Orthogonalization forward & reverse +const orth_fwd_alg = :DefaultAlgorithm +const orth_fwd_positive = true +const orth_rrule_alg = :FullPullback +const orth_rrule_verbosity = 0 # Projectors const projector_alg = :HalfInfiniteProjector # ∈ {:HalfInfiniteProjector, :FullInfiniteProjector} diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl index ebc20f091..171610fa7 100644 --- a/src/PEPSKit.jl +++ b/src/PEPSKit.jl @@ -62,8 +62,7 @@ include("utility/diffable_threads.jl") include("utility/twistdual.jl") include("utility/eigh.jl") include("utility/svd.jl") -include("utility/qr.jl") -include("utility/lq.jl") +include("utility/orth.jl") include("utility/rotations.jl") include("utility/hook_pullback.jl") include("utility/autoopt.jl") @@ -173,7 +172,7 @@ include("algorithms/select_algorithm.jl") using .Defaults: set_scheduler! export set_scheduler! -export EighAdjoint, IterEigh, SVDAdjoint, IterSVD, QRAdjoint, LQAdjoint +export EighAdjoint, IterEigh, SVDAdjoint, IterSVD, OrthAdjoint, QRAdjoint export CTMRGEnv, SequentialCTMRG, SimultaneousCTMRG export initialize_ctmrg_environment, RandomInitialization, ProductStateInitialization, ApplicationInitialization, IdentityInitialization diff --git a/src/algorithms/ctmrg/c4v.jl b/src/algorithms/ctmrg/c4v.jl index 093506ac4..c79c32e4b 100644 --- a/src/algorithms/ctmrg/c4v.jl +++ b/src/algorithms/ctmrg/c4v.jl @@ -21,7 +21,7 @@ For a full description, see [`leading_boundary`](@ref). The supported keywords a * `miniter::Int=$(Defaults.ctmrg_miniter)` * `verbosity::Int=$(Defaults.ctmrg_verbosity)` * `trunc::Union{TruncationStrategy,NamedTuple}=(; alg::Symbol=:$(Defaults.trunc))` -* `decomposition_alg::Union{NamedTuple,<:EighAdjoint,<:QRAdjoint}=(;)` +* `decomposition_alg::Union{NamedTuple,<:EighAdjoint,<:OrthAdjoint}=(;)` * `projector_alg::Symbol=:$(Defaults.projector_alg_c4v)` """ struct C4vCTMRG{P <: ProjectorAlgorithm} <: CTMRGAlgorithm diff --git a/src/algorithms/ctmrg/ctmrg.jl b/src/algorithms/ctmrg/ctmrg.jl index f1ac32a8f..b1c243c04 100644 --- a/src/algorithms/ctmrg/ctmrg.jl +++ b/src/algorithms/ctmrg/ctmrg.jl @@ -89,11 +89,9 @@ supplied via the keyword arguments or directly as an [`CTMRGAlgorithm`](@ref) st - `:FullInfiniteProjector` : Projection via SVDs of full-infinite (all four enlarged corners) CTMRG environments. - `:C4vEighProjector` : Projection via `eigh` of the Hermitian enlarged corner, works only for [`C4vCTMRG`](@ref). - `:C4vQRProjector` : Projection via QR decomposition of the lower-rank column-enlarged corner, works only for [`C4vCTMRG`](@ref). -* `decomposition_alg::Union{NamedTuple,<:SVDAdjoint,<:EighAdjoint,<:QRAdjoint}` : Tensor - decomposition algorithm used for computing projectors. When specified as a `NamedTuple`, - the settings are passed a the forward algorithm to the appropriate decomposition - for the given projector algorithm. For information on which forward algorithms are - available, and how to specify them, see [`SVDAdjoint`](@ref), [`EighAdjoint`](@ref) and [`QRAdjoint`](@ref). +* `decomposition_alg::Union{NamedTuple,<:SVDAdjoint,<:EighAdjoint,<:OrthAdjoint}` : Tensor decomposition algorithm used for computing projectors. + When specified as a `NamedTuple`, the settings are passed as the forward algorithm to the appropriate decomposition for the given projector algorithm. + For information on which forward algorithms are available, and how to specify them, see [`SVDAdjoint`](@ref), [`EighAdjoint`](@ref) and [`OrthAdjoint`](@ref). ## Return values diff --git a/src/algorithms/ctmrg/projectors/c4v_qr.jl b/src/algorithms/ctmrg/projectors/c4v_qr.jl index 987adb4aa..9ac50bd2d 100644 --- a/src/algorithms/ctmrg/projectors/c4v_qr.jl +++ b/src/algorithms/ctmrg/projectors/c4v_qr.jl @@ -13,7 +13,7 @@ $(TYPEDFIELDS) Construct the C₄ᵥ `qr`-based projector algorithm based on the following keyword arguments: -* `decomposition_alg=QRAdjoint()` : `left_orth` algorithm including the reverse rule. See [`QRAdjoint`](@ref). +* `decomposition_alg=OrthAdjoint()` : `left_orth` algorithm including the reverse rule. See [`OrthAdjoint`](@ref). """ struct C4vQRProjector{S} <: ProjectorAlgorithm # TODO: support all `left_orth` algorithms diff --git a/src/algorithms/ctmrg/projectors/projectors.jl b/src/algorithms/ctmrg/projectors/projectors.jl index 6815669f8..8d1f6c637 100644 --- a/src/algorithms/ctmrg/projectors/projectors.jl +++ b/src/algorithms/ctmrg/projectors/projectors.jl @@ -30,7 +30,7 @@ function ProjectorAlgorithm(; elseif alg in [:C4vEighProjector] _alg_or_nt(EighAdjoint, decomposition_alg) elseif alg in [:C4vQRProjector] - _alg_or_nt(QRAdjoint, decomposition_alg) + _alg_or_nt(OrthAdjoint, decomposition_alg) end # TODO: how do we solve this in a proper way? # qr-ctmrg does not need truncation or degeneracy checks diff --git a/src/utility/lq.jl b/src/utility/lq.jl deleted file mode 100644 index 37396c194..000000000 --- a/src/utility/lq.jl +++ /dev/null @@ -1,39 +0,0 @@ -""" -$(TYPEDEF) - -Wrapper for an LQ decomposition algorithm with a defined reverse rule, used by `right_orth(!)`. -See [`QRAdjoint`](@ref) for details of the constructor options and defaults. -""" -struct LQAdjoint{F, R} - fwd_alg::F - rrule_alg::R -end - -LQAdjoint(; kwargs...) = LQAdjoint(_orthogonalization_algorithms(; kwargs...)...) - -"""Compute an LQ decomposition using the forward algorithm and reverse rule configured by `LQAdjoint`.""" -MatrixAlgebraKit.right_orth(t, alg::LQAdjoint) = right_orth!(copy(t), alg) -"""Factorize the input in place using the LQ forward algorithm configured by `LQAdjoint`.""" -MatrixAlgebraKit.right_orth!(t, alg::LQAdjoint) = right_orth!(t, alg.fwd_alg) - -# right_orth! rrule wrapping MatrixAlgebraKit's lq_pullback! -function ChainRulesCore.rrule( - ::typeof(right_orth!), - t::AbstractTensorMap, - alg::LQAdjoint{F, R}, - ) where {F <: MatrixAlgebraKit.Algorithm, R <: FullPullback} - LQ = right_orth(t, alg) - gtol = _get_pullback_gauge_tol(alg.rrule_alg.verbosity) - - """Propagate LQ cotangents with the configured gauge-sensitivity tolerance.""" - function right_orth!_pullback(ΔLQ) - Δt = zeros(scalartype(t), space(t)) - MatrixAlgebraKit.lq_pullback!(Δt, t, LQ, unthunk.(ΔLQ); gauge_atol = gtol(ΔLQ)) - return NoTangent(), Δt, NoTangent() - end - function right_orth!_pullback(::Tuple{ZeroTangent, ZeroTangent}) - return NoTangent(), ZeroTangent(), NoTangent() - end - - return LQ, right_orth!_pullback -end diff --git a/src/utility/orth.jl b/src/utility/orth.jl new file mode 100644 index 000000000..db8cfae95 --- /dev/null +++ b/src/utility/orth.jl @@ -0,0 +1,139 @@ +""" +$(TYPEDEF) + +Wrapper for an orthogonalization algorithm `fwd_alg` with a defined reverse rule `rrule_alg`. +Use the same instance with `left_orth(!)` for QR decomposition and `right_orth(!)` for LQ decomposition. + +## Fields + +$(TYPEDFIELDS) + +## Constructors + + OrthAdjoint(; kwargs...) + +Construct an `OrthAdjoint` algorithm struct based on the following keyword arguments: + +* `fwd_alg::Union{Algorithm,NamedTuple}=(; alg::Symbol=$(Defaults.orth_fwd_alg))`: Orthogonalization algorithm of the forward pass which can either be passed as an `Algorithm` instance or a `NamedTuple` where the algorithm is specified by the `alg` keyword. + The available algorithms are provided through MatrixAlgebraKit and include: + - `:DefaultAlgorithm` : MatrixAlgebraKit's [default orthogonalization algorithm](@extref MatrixAlgebraKit.DefaultAlgorithm) for a given matrix type. + - `:Householder` : MatrixAlgebraKit's [`Householder`](@extref MatrixAlgebraKit.Householder) +* `rrule_alg::Union{Algorithm,NamedTuple}=(; alg::Symbol=$(Defaults.orth_rrule_alg))`: Reverse-rule algorithm for differentiating the QR or LQ decomposition. + Can be supplied by an `Algorithm` instance directly or as a `NamedTuple` where `alg` is one of the following: + - `:FullPullback` : MatrixAlgebraKit's [`qr_pullback!`](@extref MatrixAlgebraKit.qr_pullback!) or [`lq_pullback!`](@extref MatrixAlgebraKit.lq_pullback!), selected by the orthogonalization direction. + +!!! note + Manually specifying a `rrule_alg` is considered expert-mode usage, and should only be done when full control over the implementation is desired. + For all regular use cases, the default reverse rule algorithms, automatically chosen based on the forward algorithm, should be sufficient. +""" +struct OrthAdjoint{F, R} + fwd_alg::F + rrule_alg::R +end + +"""Compatibility alias for [`OrthAdjoint`](@ref), supporting both QR and LQ decompositions.""" +const QRAdjoint = OrthAdjoint + +const ORTH_FWD_SYMBOLS = IdDict{Symbol, Any}( + :DefaultAlgorithm => DefaultAlgorithm, + :Householder => Householder, +) +const ORTH_RRULE_SYMBOLS = IdDict{Symbol, Type{<:Any}}( + :FullPullback => FullPullback +) + +function OrthAdjoint(; fwd_alg = (;), rrule_alg = (;)) + # parse forward algorithm + fwd_algorithm = if fwd_alg isa NamedTuple + fwd_kwargs = (; + alg = Defaults.orth_fwd_alg, + positive = Defaults.orth_fwd_positive, + fwd_alg..., + ) # overwrite with specified kwargs + haskey(ORTH_FWD_SYMBOLS, fwd_kwargs.alg) || + throw(ArgumentError("unknown forward algorithm: $(fwd_kwargs.alg)")) + fwd_type = ORTH_FWD_SYMBOLS[fwd_kwargs.alg] + fwd_kwargs = Base.structdiff(fwd_kwargs, (; alg = nothing)) # remove `alg` keyword argument + fwd_type(; fwd_kwargs...) + else + fwd_alg + end + + # parse reverse-rule algorithm + rrule_algorithm = if rrule_alg isa NamedTuple + rrule_kwargs = (; + alg = Defaults.orth_rrule_alg, + verbosity = Defaults.orth_rrule_verbosity, + rrule_alg..., + ) # overwrite with specified kwargs + + haskey(ORTH_RRULE_SYMBOLS, rrule_kwargs.alg) || + throw(ArgumentError("unknown rrule algorithm: $(rrule_kwargs.alg)")) + rrule_type = ORTH_RRULE_SYMBOLS[rrule_kwargs.alg] + rrule_type() + else + rrule_alg + end + + return OrthAdjoint(fwd_algorithm, rrule_algorithm) +end + +""" + left_orth(t, alg::OrthAdjoint) + left_orth!(t, alg::OrthAdjoint) + +Wrapper around `left_orth(!)` which dispatches on the `OrthAdjoint` algorithm. +This is needed since a custom adjoint may be defined, depending on the `alg`. +""" +MatrixAlgebraKit.left_orth(t, alg::OrthAdjoint) = left_orth!(copy(t), alg) +MatrixAlgebraKit.left_orth!(t, alg::OrthAdjoint) = left_orth!(t, alg.fwd_alg) + +# left_orth! rrule wrapping MatrixAlgebraKit's qr_pullback! +# https://github.com/QuantumKitHub/MatrixAlgebraKit.jl/blob/b76c7bb60014ecfead6925d0df6cb4b8d7c2668a/src/pullbacks/qr.jl#L49 +function ChainRulesCore.rrule( + ::typeof(left_orth!), + t::AbstractTensorMap, + alg::OrthAdjoint{F, R}, + ) where {F <: MatrixAlgebraKit.Algorithm, R <: FullPullback} + QR = left_orth(t, alg) + gtol = _get_pullback_gauge_tol(alg.rrule_alg.verbosity) + + """Propagate QR cotangents with the configured gauge-sensitivity tolerance.""" + function left_orth!_pullback(ΔQR) + Δt = zeros(storagetype(t), space(t)) + MatrixAlgebraKit.qr_pullback!(Δt, t, QR, unthunk.(ΔQR); gauge_atol = gtol(ΔQR)) + return NoTangent(), Δt, NoTangent() + end + function left_orth!_pullback(::Tuple{ZeroTangent, ZeroTangent}) + return NoTangent(), ZeroTangent(), NoTangent() + end + + return QR, left_orth!_pullback +end + +"""Compute an LQ decomposition using the forward algorithm and reverse rule configured by `OrthAdjoint`.""" +MatrixAlgebraKit.right_orth(t, alg::OrthAdjoint) = right_orth!(copy(t), alg) +"""Factorize the input in place using the LQ forward algorithm configured by `OrthAdjoint`.""" +MatrixAlgebraKit.right_orth!(t, alg::OrthAdjoint) = right_orth!(t, alg.fwd_alg) + +# right_orth! rrule wrapping MatrixAlgebraKit's lq_pullback! +function ChainRulesCore.rrule( + ::typeof(right_orth!), + t::AbstractTensorMap, + alg::OrthAdjoint{F, R}, + ) where {F <: MatrixAlgebraKit.Algorithm, R <: FullPullback} + LQ = right_orth(t, alg) + gtol = _get_pullback_gauge_tol(alg.rrule_alg.verbosity) + + """Propagate LQ cotangents with the configured gauge-sensitivity tolerance.""" + function right_orth!_pullback(ΔLQ) + Δt = zeros(storagetype(t), space(t)) + MatrixAlgebraKit.lq_pullback!(Δt, t, LQ, unthunk.(ΔLQ); gauge_atol = gtol(ΔLQ)) + return NoTangent(), Δt, NoTangent() + end + function right_orth!_pullback(::Tuple{ZeroTangent, ZeroTangent}) + return NoTangent(), ZeroTangent(), NoTangent() + end + + return LQ, right_orth!_pullback +end diff --git a/src/utility/qr.jl b/src/utility/qr.jl deleted file mode 100644 index 372f075f0..000000000 --- a/src/utility/qr.jl +++ /dev/null @@ -1,112 +0,0 @@ -""" -$(TYPEDEF) - -Wrapper for a QR decomposition algorithm `fwd_alg` with a defined reverse rule `rrule_alg`. - -## Fields - -$(TYPEDFIELDS) - -## Constructors - - QRAdjoint(; kwargs...) - -Construct a `QRAdjoint` algorithm struct based on the following keyword arguments: - -* `fwd_alg::Union{Algorithm,NamedTuple}=(; alg::Symbol=$(Defaults.qr_fwd_alg))`: QR - algorithm of the forward pass which can either be passed as an `Algorithm` instance or a - `NamedTuple` where the algorithm is specified by the `alg` keyword. - The available algorithms are provided through MatrixAlgebraKit and include: - - `:DefaultAlgorithm` : MatrixAlgebraKit's [default QR algorithm](@extref MatrixAlgebraKit.DefaultAlgorithm) for a given matrix type. - - `:Householder` : MatrixAlgebraKit's [`Householder`](@extref MatrixAlgebraKit.Householder) -* `rrule_alg::Union{Algorithm,NamedTuple}=(; alg::Symbol=$(Defaults.qr_rrule_alg))`: Reverse-rule algorithm for differentiating the QR decomposition. - Can be supplied by an `Algorithm` instance directly or as a `NamedTuple` where `alg` is one of the following: - - `:FullPullback` : MatrixAlgebraKit's [`qr_pullback!`](@extref MatrixAlgebraKit.qr_pullback!) - -!!! note - Manually specifying a `rrule_alg` is considered expert-mode usage, and should only be done when full control over the implementation is desired. - For all regular use cases, the default reverse rule algorithms, automatically chosen based on the forward algorithm, should be sufficient. -""" -struct QRAdjoint{F, R} - fwd_alg::F - rrule_alg::R -end - -const QR_FWD_SYMBOLS = IdDict{Symbol, Any}( - :DefaultAlgorithm => DefaultAlgorithm, - :Householder => Householder, -) -const QR_RRULE_SYMBOLS = IdDict{Symbol, Type{<:Any}}( - :FullPullback => FullPullback -) - -QRAdjoint(; kwargs...) = QRAdjoint(_orthogonalization_algorithms(; kwargs...)...) - -"""Parse the shared forward and reverse algorithm options for QR and LQ decompositions.""" -function _orthogonalization_algorithms(; fwd_alg = (;), rrule_alg = (;)) - # parse forward algorithm - fwd_algorithm = if fwd_alg isa NamedTuple - fwd_kwargs = (; - alg = Defaults.qr_fwd_alg, - positive = Defaults.qr_fwd_positive, - fwd_alg..., - ) # overwrite with specified kwargs - haskey(QR_FWD_SYMBOLS, fwd_kwargs.alg) || - throw(ArgumentError("unknown forward algorithm: $(fwd_kwargs.alg)")) - fwd_type = QR_FWD_SYMBOLS[fwd_kwargs.alg] - fwd_kwargs = Base.structdiff(fwd_kwargs, (; alg = nothing)) # remove `alg` keyword argument - fwd_type(; fwd_kwargs...) - else - fwd_alg - end - - # parse reverse-rule algorithm - rrule_algorithm = if rrule_alg isa NamedTuple - rrule_kwargs = (; - alg = Defaults.qr_rrule_alg, - verbosity = Defaults.qr_rrule_verbosity, - rrule_alg..., - ) # overwrite with specified kwargs - - haskey(QR_RRULE_SYMBOLS, rrule_kwargs.alg) || - throw(ArgumentError("unknown rrule algorithm: $(rrule_kwargs.alg)")) - rrule_type = QR_RRULE_SYMBOLS[rrule_kwargs.alg] - rrule_type() - else - rrule_alg - end - - return fwd_algorithm, rrule_algorithm -end - -""" - left_orth(t, alg::QRAdjoint) - left_orth!(t, alg::QRAdjoint) - -Wrapper around `left_orth(!)` which dispatches on the `QRAdjoint` algorithm. -This is needed since a custom adjoint may be defined, depending on the `alg`. -""" -MatrixAlgebraKit.left_orth(t, alg::QRAdjoint) = left_orth!(copy(t), alg) -MatrixAlgebraKit.left_orth!(t, alg::QRAdjoint) = left_orth!(t, alg.fwd_alg) - -# left_orth! rrule wrapping MatrixAlgebraKit's qr_pullback! -# https://github.com/QuantumKitHub/MatrixAlgebraKit.jl/blob/b76c7bb60014ecfead6925d0df6cb4b8d7c2668a/src/pullbacks/qr.jl#L49 -function ChainRulesCore.rrule( - ::typeof(left_orth!), - t::AbstractTensorMap, - alg::QRAdjoint{F, R}, - ) where {F <: MatrixAlgebraKit.Algorithm, R <: FullPullback} - QR = left_orth(t, alg) - gtol = _get_pullback_gauge_tol(alg.rrule_alg.verbosity) - - function left_orth!_pullback(ΔQR) - Δt = zeros(storagetype(t), space(t)) - MatrixAlgebraKit.qr_pullback!(Δt, t, QR, unthunk.(ΔQR); gauge_atol = gtol(ΔQR)) - return NoTangent(), Δt, NoTangent() - end - function left_orth!_pullback(::Tuple{ZeroTangent, ZeroTangent}) - return NoTangent(), ZeroTangent(), NoTangent() - end - - return QR, left_orth!_pullback -end diff --git a/test/testsuite/TestSuite.jl b/test/testsuite/TestSuite.jl index fd4223c5b..46040ef3c 100644 --- a/test/testsuite/TestSuite.jl +++ b/test/testsuite/TestSuite.jl @@ -266,11 +266,11 @@ module UtilityEighWrapper end using .UtilityEighWrapper -module UtilityQRWrapper - include("utility/qr_wrapper.jl") - export utility_qr_wrapper +module UtilityOrthWrapper + include("utility/orth_wrapper.jl") + export utility_orth_wrapper end -using .UtilityQRWrapper +using .UtilityOrthWrapper module UtilitySVDWrapper include("utility/svd_wrapper.jl") diff --git a/test/testsuite/gradients/c4v_ctmrg_gradients.jl b/test/testsuite/gradients/c4v_ctmrg_gradients.jl index 96dd1f6c8..91113aec2 100644 --- a/test/testsuite/gradients/c4v_ctmrg_gradients.jl +++ b/test/testsuite/gradients/c4v_ctmrg_gradients.jl @@ -29,7 +29,7 @@ projector_algs = (:C4vEighProjector, :C4vQRProjector) decomposition_alg(projector_alg, rrule_alg) = if projector_alg == :C4vEighProjector EighAdjoint(; rrule_alg = (; alg = rrule_alg)) elseif projector_alg == :C4vQRProjector - QRAdjoint(; rrule_alg = (; alg = rrule_alg)) + OrthAdjoint(; rrule_alg = (; alg = rrule_alg)) else error("unknown projector alg: $projector_alg") end diff --git a/test/testsuite/utility/orth_wrapper.jl b/test/testsuite/utility/orth_wrapper.jl new file mode 100644 index 000000000..7b416cade --- /dev/null +++ b/test/testsuite/utility/orth_wrapper.jl @@ -0,0 +1,56 @@ +using Test, Random, LinearAlgebra +using TensorKit, MatrixAlgebraKit, PEPSKit, Zygote, Adapt +using ChainRulesCore: rrule, NoTangent, ZeroTangent + +"""Check shared QR/LQ algorithms and pullbacks, including mixed dual fermionic legs.""" +function utility_orth_wrapper(AT) + Random.seed!(73194) + algorithms = (OrthAdjoint(), OrthAdjoint(; fwd_alg = (; alg = :Householder))) + spaces = (ComplexSpace(2), Vect[FermionParity](0 => 1, 1 => 1), Vect[FermionNumber](-1 => 1, 0 => 1, 1 => 1)) + return @testset "Orthogonalization wrapper ($AT)" begin + @test QRAdjoint === OrthAdjoint + @test QRAdjoint() isa OrthAdjoint + fwd_alg = Householder(; positive = true) + rrule_alg = PEPSKit.FullPullback(; verbosity = 1) + configured = OrthAdjoint(; fwd_alg, rrule_alg) + @test configured.fwd_alg === fwd_alg + @test configured.rrule_alg === rrule_alg + @test configured == OrthAdjoint(fwd_alg, rrule_alg) + + @testset "$(typeof(alg.fwd_alg)), $(sectortype(V)), $nout ← $nin, $T" for alg in algorithms, V in spaces, (nout, nin) in ((2, 3), (3, 2)), T in (Float64, ComplexF64) + output, input = (prod(ntuple(i -> isodd(i) ? V : V', n)) for n in (nout, nin)) + A = adapt(AT, randn(T, output ← input)) + H, K = (adapt(AT, randn(T, W ← W)) for W in (output, input)) + direction = normalize(adapt(AT, randn(T, space(A)))) + + @testset "$orth" for (orth, orth!, opposite) in ((left_orth, left_orth!, right_orth), (right_orth, right_orth!, left_orth)) + factorize = X -> orth(X, alg) + # Explicit tuple adjoints avoid costly inference in the broadcast pullback. + reference = X -> let (a, b) = opposite(X', alg) + (b', a') + end + A_before = copy(A) + a, b = factorize(A) + @test A == A_before + @test all(isapprox.((a, b), reference(A))) + @test a * b ≈ A + @test orth!(copy(A), alg) == (a, b) + + # The real Gram-matrix loss is gauge invariant and exercises both factors. + loss = decomposition -> X -> let (a, b) = decomposition(X) + real(dot(H, a * a') + dot(K, b' * b)) + end + f, fref = loss(factorize), loss(reference) + gradient = only(Zygote.gradient(f, A)) + @test gradient ≈ only(Zygote.gradient(fref, A)) + @test storagetype(gradient) == storagetype(A) + epsilon = 1.0e-5 + finite = (f(A + epsilon * direction) - f(A - epsilon * direction)) / (2 * epsilon) + @test real(dot(gradient, direction)) ≈ finite atol = 1.0e-7 rtol = 1.0e-7 + + _, back = rrule(orth!, copy(A), alg) + @test back((ZeroTangent(), ZeroTangent())) == (NoTangent(), ZeroTangent(), NoTangent()) + end + end + end +end diff --git a/test/testsuite/utility/qr_wrapper.jl b/test/testsuite/utility/qr_wrapper.jl deleted file mode 100644 index 348cc66ce..000000000 --- a/test/testsuite/utility/qr_wrapper.jl +++ /dev/null @@ -1,39 +0,0 @@ -using Test, Random, LinearAlgebra -using TensorKit, MatrixAlgebraKit, PEPSKit, Zygote, Adapt -using ChainRulesCore: rrule, NoTangent, ZeroTangent - -"""Check LQ factors and pullbacks against adjointed QR, including mixed dual fermionic legs.""" -function utility_qr_wrapper(AT) - Random.seed!(73194) - qr_alg, lq_alg = QRAdjoint(), LQAdjoint() - factorize = X -> right_orth(X, lq_alg) - # Explicit tuple adjoints avoid costly inference in the broadcast pullback. - reference = X -> let (q, r) = left_orth(X', qr_alg) - (r', q') - end - spaces = (ComplexSpace(2), Vect[FermionParity](0 => 1, 1 => 1), Vect[FermionNumber](-1 => 1, 0 => 1, 1 => 1)) - return @testset "QR/LQ wrapper ($(sectortype(V)), $nout ← $nin, $AT)" for V in spaces, (nout, nin) in ((2, 3), (3, 2)) - output, input = (prod(ntuple(i -> isodd(i) ? V : V', n)) for n in (nout, nin)) - A = adapt(AT, randn(ComplexF64, output ← input)) - L, Q = factorize(A) - @test all(isapprox.((L, Q), reference(A))) - @test L * Q ≈ A - @test right_orth!(copy(A), lq_alg) == (L, Q) - - H, K = (adapt(AT, randn(ComplexF64, W ← W)) for W in (output, input)) - # The real Gram-matrix loss is gauge invariant and exercises both factors. - loss = decomposition -> X -> let (l, q) = decomposition(X) - real(dot(H, l * l') + dot(K, q' * q)) - end - f, fref = loss(factorize), loss(reference) - gradient = only(Zygote.gradient(f, A)) - @test gradient ≈ only(Zygote.gradient(fref, A)) - direction = normalize(adapt(AT, randn(ComplexF64, space(A)))) - epsilon = 1.0e-5 - finite = (f(A + epsilon * direction) - f(A - epsilon * direction)) / (2 * epsilon) - @test real(dot(gradient, direction)) ≈ finite atol = 1.0e-7 rtol = 1.0e-7 - - _, back = rrule(right_orth!, copy(A), lq_alg) - @test back((ZeroTangent(), ZeroTangent())) == (NoTangent(), ZeroTangent(), NoTangent()) - end -end diff --git a/test/utility/qr_wrapper.jl b/test/utility/orth_wrapper.jl similarity index 81% rename from test/utility/qr_wrapper.jl rename to test/utility/orth_wrapper.jl index f39bea571..9003477e6 100644 --- a/test/utility/qr_wrapper.jl +++ b/test/utility/orth_wrapper.jl @@ -7,5 +7,5 @@ using .TestSuite is_buildkite = get(ENV, "BUILDKITE", "false") == "true" if !is_buildkite - TestSuite.utility_qr_wrapper(Vector) + TestSuite.utility_orth_wrapper(Vector) end From 5455b896fe9f9fbe7d8aee848339a67a292a618f Mon Sep 17 00:00:00 2001 From: Yue Zhengyuan Date: Mon, 19 Oct 2026 21:21:52 +0800 Subject: [PATCH 4/4] Match docstring of `right_orth` Co-authored-by: Paul Brehmer --- src/utility/orth.jl | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/src/utility/orth.jl b/src/utility/orth.jl index db8cfae95..232ddeaf5 100644 --- a/src/utility/orth.jl +++ b/src/utility/orth.jl @@ -111,9 +111,14 @@ function ChainRulesCore.rrule( return QR, left_orth!_pullback end -"""Compute an LQ decomposition using the forward algorithm and reverse rule configured by `OrthAdjoint`.""" +""" + right_orth(t, alg::OrthAdjoint) + right_orth!(t, alg::OrthAdjoint) + +Wrapper around `right_orth(!)` which dispatches on the `OrthAdjoint` algorithm. +This is needed since a custom adjoint may be defined, depending on the `alg`. +""" MatrixAlgebraKit.right_orth(t, alg::OrthAdjoint) = right_orth!(copy(t), alg) -"""Factorize the input in place using the LQ forward algorithm configured by `OrthAdjoint`.""" MatrixAlgebraKit.right_orth!(t, alg::OrthAdjoint) = right_orth!(t, alg.fwd_alg) # right_orth! rrule wrapping MatrixAlgebraKit's lq_pullback!