diff --git a/src/TensorAlgebra.jl b/src/TensorAlgebra.jl index d5f03e6..b6bcc51 100644 --- a/src/TensorAlgebra.jl +++ b/src/TensorAlgebra.jl @@ -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 diff --git a/src/factorizations.jl b/src/factorizations.jl index 38ae925..758a31e 100644 --- a/src/factorizations.jl +++ b/src/factorizations.jl @@ -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...) + 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 diff --git a/test/test_exports.jl b/test/test_exports.jl index 67e6b20..cf4171c 100644 --- a/test/test_exports.jl +++ b/test/test_exports.jl @@ -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, diff --git a/test/test_factorizations.jl b/test/test_factorizations.jl index 6e022f9..080a530 100644 --- a/test/test_factorizations.jl +++ b/test/test_factorizations.jl @@ -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 @@ -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