Skip to content

Add qr_compact!, forming Q in the input's memory #243

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jack-dunham wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Add qr_compact!, forming Q in the input's memory #243

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion src/TensorAlgebra.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ export contract, contract!, contractalign, dual, isdual, MatrixAlgebra
if VERSION >= v"1.11.0-DEV.469"
eval(
Meta.parse(
"public AbstractAlgorithm, add!, AddBroadcasted, addends, allocate_output, allocate_project, arguments, axes, bipartition, bipartition_axes, biperm, bipermutedims, bipermutedims!, bipermutedimsopadd!, cat_axis, cat_similar, check_input, concatenate, concatenate!, ConjBroadcasted, contractadd!, ContractAlgorithm, contractopadd!, contractperm, contractperm!, contractpermadd!, contractpermalign, contractpermopadd!, data, datatype, default_algorithm, dims2cat, directsum, dot, dotperm, eig_full, eig_trunc, eig_vals, eigh_full, eigh_trunc, eigh_vals, fill_map, flattenlinear, has_bipartition, infer_aux_space, invsqrth_safe, is_output_view, is_projected, isidentitybiperm, label_type, left_null, left_orth, left_polar, LinearBroadcasted, linearbroadcasted, lq_compact, lq_full, matricize, MatricizeContract, matricizeop, matricizeop!, matricizeopcopy, matricizeopview, MATRIX_FUNCTIONS, ndims, ndims_codomain, ndims_domain, one, ones_map, operation, output_axes, PermutedDims, permuteddims, permutedims, permutedims!, permutedimsadd!, permutedimsop, permutedimsopadd!, project, project!, project_aux, project_hermitian, projectto!, qr_compact, qr_full, rand_map, randn_map, right_null, right_orth, right_polar, scalar, scale!, scaleadd!, ScaledBroadcasted, select_algorithm, similar_map, size, sqrth_invsqrth_safe, sqrth_safe, sum, svd_compact, svd_full, svd_trunc, svd_vals, TensorOperationsContract, to_range, tr, trivialrange, tryflattenlinear, tryproject, tryproject_aux, twist!, unchecked_project, unchecked_project_aux, ungrade, unmatricize, unmatricize!, unmatricize_factors, unmatricizeadd!, unproject, unscaled, zero!, zeros_map"
"public AbstractAlgorithm, add!, AddBroadcasted, addends, allocate_output, allocate_project, arguments, axes, bipartition, bipartition_axes, biperm, bipermutedims, bipermutedims!, bipermutedimsopadd!, cat_axis, cat_similar, check_input, concatenate, concatenate!, ConjBroadcasted, contractadd!, ContractAlgorithm, contractopadd!, contractperm, contractperm!, contractpermadd!, contractpermalign, contractpermopadd!, data, datatype, default_algorithm, dims2cat, directsum, dot, dotperm, eig_full, eig_trunc, eig_vals, eigh_full, eigh_trunc, eigh_vals, fill_map, flattenlinear, has_bipartition, infer_aux_space, invsqrth_safe, is_output_view, is_projected, isidentitybiperm, label_type, left_null, left_orth, left_polar, LinearBroadcasted, linearbroadcasted, lq_compact, lq_full, matricize, MatricizeContract, matricizeop, matricizeop!, matricizeopcopy, matricizeopview, MATRIX_FUNCTIONS, ndims, ndims_codomain, ndims_domain, one, ones_map, operation, output_axes, PermutedDims, permuteddims, permutedims, permutedims!, permutedimsadd!, permutedimsop, permutedimsopadd!, project, project!, project_aux, project_hermitian, projectto!, qr_compact, qr_compact!, qr_full, rand_map, randn_map, right_null, right_orth, right_polar, scalar, scale!, scaleadd!, ScaledBroadcasted, select_algorithm, similar_map, size, sqrth_invsqrth_safe, sqrth_safe, sum, svd_compact, svd_full, svd_trunc, svd_vals, TensorOperationsContract, to_range, tr, trivialrange, tryflattenlinear, tryproject, tryproject_aux, twist!, unchecked_project, unchecked_project_aux, ungrade, unmatricize, unmatricize!, unmatricize_factors, unmatricizeadd!, unproject, unscaled, zero!, zeros_map"
)
)
end
Expand Down
40 changes: 40 additions & 0 deletions src/factorizations.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -162,6 +162,46 @@ See also `MatrixAlgebraKit.qr_compact!`.
"""
qr_compact

"""
qr_compact!(A, labels_A, labels_codomain, labels_domain; kwargs...) -> Q, R
qr_compact!(A, perm_codomain, perm_domain; kwargs...) -> Q, R
qr_compact!(A, ndims_codomain::Val; kwargs...) -> Q, R

Like [`qr_compact`](@ref), but `A` may be overwritten. When `A` is dense and its matricization
has no more columns than rows, `Q` is formed in `A`'s memory.
"""
function qr_compact!(A, perm_codomain, perm_domain; kwargs...)
ndims(A) == length(perm_codomain) + length(perm_domain) ||
throw(ArgumentError("Invalid bipermutation"))

m = prod(i -> size(A, i), perm_codomain; init = 1)
n = prod(i -> size(A, i), perm_domain; init = 1)
if A isa DenseArray && m >= n
A_mat = reshape(A, (m, n))
# The permuted matricization is a temporary; copying it back puts `Q` in `A`'s memory.
if !is_output_view(matricizeop, identity, A, perm_codomain, perm_domain)
copyto!(A_mat, matricizeopcopy(identity, A, perm_codomain, perm_domain))
end
F = MatrixAlgebraKit.qr_compact!(A_mat, (A_mat, similar(A_mat, (n, n))); kwargs...)
Comment on lines +180 to +185

@mtfishman mtfishman Oct 13, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we need this branch? I.e. do we really care if A is overwritten? I think we can just have a "contract" that A might be overwritten, and you can't rely on it any more, and Q should be obtained from the return value.

@jack-dunham jack-dunham Oct 13, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We do care that A is strictly overwritten for the following reason:

It allows the intermediate (generated by matricizeopcopy) to be garbage collected. For context, this function is to be used when one really wants to reuse the memory of A to avoid having to manually invoke the garbage collector on it, e.g. in cases where one cannot guarantee cases where all references to A are removed. An example is where A = state[vertex], and one wants to reuse the memory of A immediately.

In more detail: If one must permute A to get it the correct form for QR, then one must allocate a temporary C. If one calls QR on this temporary C directly (writing Q into this array), then there are now two A sized tensors in memory: Q (written into C memory), and the input A. That is fine and unavoidable.

Now, if one does not copy this permuted form back into A before the QR, then this Q returned is additional allocated memory to A now existing outside the scope qr_compact!. In addition, the tensor A exists until it is garbage collected which will not occur when the reference state[vertex] still exists.

If we instead overwrite A (which is allowed by the ! contract) using the copyto!, and QR this, then the intermediate C can get garbage collected when qr_compact! returns, and one is left with only one A sized tensor allocated in the calling scope.

One could remove the if statement and have the body of the if statement apply in all cases, but then one gets an unnecessary allocation and call to copyto! when the indices are already in the correct order.

We can change the contract to might be overwritten, but then this function is not useful to me in the one and only place it's used in the gate application code, and I am not really sure what the advantage of that would be.

else
F = MatrixAlgebraKit.qr_compact!(matricize(A, perm_codomain, perm_domain); kwargs...)
end

axes_codomain, axes_domain = bipartition_axes(
map(i -> axes(A, i), (perm_codomain..., perm_domain...)),
Val(length(perm_codomain))
)
return unmatricize_factors(qr_compact, F, axes_codomain, axes_domain)
end
function qr_compact!(A, ndims_codomain::Val; kwargs...)
return qr_compact!(A, identitybiperm(ndims_codomain, Val(ndims(A)))...; kwargs...)
end
function qr_compact!(A, labels_A, labels_codomain, labels_domain; kwargs...)
perm_codomain, perm_domain =
biperm(Tuple.((labels_A, labels_codomain, labels_domain))...)
return qr_compact!(A, perm_codomain, perm_domain; kwargs...)
end

"""
qr_full(A, labels_A, labels_codomain, labels_domain; kwargs...) -> Q, R
qr_full(A, perm_codomain, perm_domain; kwargs...) -> Q, R
Expand Down
2 changes: 1 addition & 1 deletion test/test_exports.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ using Test: @test, @testset
:output_axes,
:PermutedDims, :permuteddims, :permutedims, :permutedims!,
:permutedimsadd!, :permutedimsop, :permutedimsopadd!, :project,
:project!, :project_aux, :project_hermitian, :projectto!, :qr_compact,
:project!, :project_aux, :project_hermitian, :projectto!, :qr_compact, :qr_compact!,
:qr_full, :rand_map, :randn_map, :right_null, :right_orth, :right_polar,
:scalar, :scale!, :scaleadd!, :ScaledBroadcasted, :select_algorithm,
:similar_map,
Expand Down
28 changes: 26 additions & 2 deletions test/test_factorizations.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,8 +1,9 @@
using LinearAlgebra: LinearAlgebra, Diagonal, I, diag, norm
using MatrixAlgebraKit: truncrank
using TensorAlgebra: TensorAlgebra, contract, contractalign, eig_full, eig_vals, eigh_full,
eigh_vals, left_null, left_orth, left_polar, lq_compact, lq_full, qr_compact, qr_full,
right_null, right_orth, right_polar, svd_compact, svd_full, svd_trunc, svd_vals
eigh_vals, left_null, left_orth, left_polar, lq_compact, lq_full, qr_compact,
qr_compact!, qr_full, right_null, right_orth, right_polar, svd_compact, svd_full,
svd_trunc, svd_vals
using Test: @test, @testset
using TestExtras: @constinferred

Expand Down Expand Up @@ -47,6 +48,29 @@ end
@test size(Q, 3) == min(size(A, 1) * size(A, 2), size(A, 3) * size(A, 4))
end

@testset "In-place compact QR ($T)" for T in elts
A = randn(T, 4, 5, 2, 3)
labels_A = (:a, :b, :c, :d)

# Tall: `Q` is formed in `A`'s memory, with or without a permutation.
Q_ref, R_ref = qr_compact(A, Val(2))
A′ = copy(A)
Q, R = @constinferred qr_compact!(A′, Val(2))
@test Q ≈ Q_ref && R ≈ R_ref
@test pointer(Q) == pointer(A′)

Q_ref, R_ref = qr_compact(A, labels_A, (:b, :d), (:c, :a))
A′ = copy(A)
Q, R = qr_compact!(A′, labels_A, (:b, :d), (:c, :a))
@test Q ≈ Q_ref && R ≈ R_ref
@test pointer(Q) == pointer(A′)

# Wide: `Q` is allocated.
Q_ref, R_ref = qr_compact(A, Val(1))
Q, R = qr_compact!(copy(A), Val(1))
@test Q ≈ Q_ref && R ≈ R_ref
end

# LQ Decomposition
# ----------------
@testset "Full LQ ($T)" for T in elts
Expand Down
Loading