From 6c3ecb46dd492003811e131a869fdfacea56f46e Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Fri, 21 Aug 2026 23:04:22 -0400 Subject: [PATCH 1/5] Matricize factorization inputs lazily and donate owned buffers --- src/factorizations.jl | 249 ++++++++++++++++++++++------------- src/matricize.jl | 26 +++- src/matrixfunctions.jl | 31 ++++- test/test_factorizations.jl | 42 ++++++ test/test_matrixfunctions.jl | 9 ++ 5 files changed, 260 insertions(+), 97 deletions(-) diff --git a/src/factorizations.jl b/src/factorizations.jl index 11a38a4d..01c2383f 100644 --- a/src/factorizations.jl +++ b/src/factorizations.jl @@ -9,50 +9,70 @@ using MatrixAlgebraKit: MatrixAlgebraKit # bond is dualized to codomain-facing form (`conj`, a no-op on a dense axis) when it lands on the # domain side of the reconstruction, matching the `unmatricize`/`similar_map` axis convention. -# Two-output factorizations: the first factor `X` has the codomain axes plus a trailing -# rank axis, the second factor `Y` has a leading rank axis plus the domain axes. +# Whether the matricized form `A_mat` is detached from `A`: fresh storage the wrapper owns, so +# a mutating consumer cannot corrupt the caller's data. Aliasing is only tracked (through +# `Base.dataids`) for `AbstractArray`s; any other tensor type (such as a `TensorMap`, whose +# trivial `permute` returns the input itself) is conservatively treated as attached. +isdetached(A_mat::AbstractArray, A::AbstractArray) = !Base.mightalias(A_mat, A) +isdetached(A_mat, A) = false + +# Skip the copy inside the non-mutating MatrixAlgebraKit entry (`f(A) = f!(copy_input(f, A))`) +# when the wrapper owns `A_mat` and it already has the floating-point eltype `copy_input` +# would produce. When `A_mat` may alias the caller's input (a dense reshape or a stored graded +# matrix), the copying entry is required: mutating `A_mat` would corrupt that input. +function maybe_donate(f, f!, A_mat, owned::Bool; kwargs...) + owned && eltype(A_mat) === float(eltype(A_mat)) && return f!(A_mat; kwargs...) + return f(A_mat; kwargs...) +end + +# `matricized(f, style, A_mat, owned, axes_codomain, axes_domain; kwargs...)` is the shared +# matrix-level body of the wrapper `f`: apply the matrix-level `f` to the matricized input +# `A_mat` and unfold the outputs with the bipartitioned axes (in the `unmatricize` convention, +# domain axes un-dualized). `owned` (see `isdetached`) gates donating `A_mat` to the mutating +# matrix-level entry. The wrappers produce `A_mat` with the maybe-alias `matricize`/ +# `matricize_input` spelling, so the permuted forms skip the eager `bipermutedims` copy. for f in ( :qr_compact, :qr_full, :lq_compact, :lq_full, :left_polar, :right_polar, :left_orth, :right_orth, + :svd_compact, :svd_full, :svd_trunc, :svd_vals, + :eigh_full, :eig_full, :eigh_trunc, :eig_trunc, :eigh_vals, :eig_vals, + :left_null, :right_null, :gram_eigh_full, :gram_eigh_full_with_pinv, + :sqrth_safe, :invsqrth_safe, :sqrth_invsqrth_safe, :project_hermitian, ) @eval begin function $f(style::MatricizeStyle, A, ndims_codomain::Val; kwargs...) A_mat = matricize(style, A, ndims_codomain) - X, Y = MatrixAlgebraKit.$f(A_mat; kwargs...) axes_codomain, axes_domain = bipartition_axes(axes(A), ndims_codomain) - return unmatricize(style, X, axes_codomain, (conj(axes(X, ndims(X))),)), - unmatricize(style, Y, (axes(Y, 1),), axes_domain) + return matricized( + $f, style, A_mat, isdetached(A_mat, A), + axes_codomain, axes_domain; kwargs... + ) end function $f(A, ndims_codomain::Val; kwargs...) return $f(MatricizeStyle(A), A, ndims_codomain; kwargs...) end - end -end -for f in ( - :qr_compact, :qr_full, :lq_compact, :lq_full, - :left_polar, :right_polar, :left_orth, :right_orth, - :svd_compact, :svd_full, :svd_trunc, :svd_vals, - :eigh_full, :eig_full, :eigh_trunc, :eig_trunc, :eigh_vals, :eig_vals, - :left_null, :right_null, :gram_eigh_full, :gram_eigh_full_with_pinv, :one, - :sqrth_safe, :invsqrth_safe, :sqrth_invsqrth_safe, :project_hermitian, - ) - @eval begin function $f( style::MatricizeStyle, A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs... ) - A_perm = bipermutedims(A, perm_codomain, perm_domain) - return $f(style, A_perm, Val(length(perm_codomain)); kwargs...) + A_mat = matricize_input(style, A, perm_codomain, perm_domain) + axes_codomain, axes_domain = bipartition_axes( + map(i -> axes(A, i), (perm_codomain..., perm_domain...)), + Val(length(perm_codomain)) + ) + return matricized( + $f, style, A_mat, isdetached(A_mat, A), + axes_codomain, axes_domain; kwargs... + ) end function $f( A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs... ) - A_perm = bipermutedims(A, perm_codomain, perm_domain) - return $f(A_perm, Val(length(perm_codomain)); kwargs...) + return $f(MatricizeStyle(A), A, perm_codomain, perm_domain; kwargs...) end function $f( @@ -71,6 +91,27 @@ for f in ( end end +# Two-output factorizations: the first factor `X` has the codomain axes plus a trailing +# rank axis, the second factor `Y` has a leading rank axis plus the domain axes. +for f in ( + :qr_compact, :qr_full, :lq_compact, :lq_full, + :left_polar, :right_polar, :left_orth, :right_orth, + ) + @eval begin + function matricized( + ::typeof($f), style::MatricizeStyle, A_mat, owned::Bool, + axes_codomain, axes_domain; kwargs... + ) + X, Y = maybe_donate( + MatrixAlgebraKit.$f, MatrixAlgebraKit.$(Symbol(f, :!)), A_mat, owned; + kwargs... + ) + return unmatricize(style, X, axes_codomain, (conj(axes(X, ndims(X))),)), + unmatricize(style, Y, (axes(Y, 1),), axes_domain) + end + end +end + """ TensorAlgebra.tr(A, labels_A, labels_codomain, labels_domain) TensorAlgebra.tr(A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}) @@ -104,8 +145,7 @@ function tr(A, ndims_codomain::Val) return tr(MatricizeStyle(A), A, ndims_codomain) end function tr(A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}) - A_perm = bipermutedims(A, perm_codomain, perm_domain) - return tr(A_perm, Val(length(perm_codomain))) + return LinearAlgebra.tr(matricizeperm(A, perm_codomain, perm_domain)) end function tr(A, labels_A, labels_codomain, labels_domain) perm_codomain, perm_domain = @@ -257,64 +297,65 @@ right_orth # rank × rank spectrum, and `Vᴴ` carries a leading rank axis plus the domain axes. for f in (:svd_compact, :svd_full) @eval begin - function $f(style::MatricizeStyle, A, ndims_codomain::Val; kwargs...) - A_mat = matricize(style, A, ndims_codomain) - U, S, Vᴴ = MatrixAlgebraKit.$f(A_mat; kwargs...) - axes_codomain, axes_domain = bipartition_axes(axes(A), ndims_codomain) + function matricized( + ::typeof($f), style::MatricizeStyle, A_mat, owned::Bool, + axes_codomain, axes_domain; kwargs... + ) + U, S, Vᴴ = maybe_donate( + MatrixAlgebraKit.$f, MatrixAlgebraKit.$(Symbol(f, :!)), A_mat, owned; + kwargs... + ) return unmatricize(style, U, axes_codomain, (conj(axes(U, ndims(U))),)), S, unmatricize(style, Vᴴ, (axes(Vᴴ, 1),), axes_domain) end - function $f(A, ndims_codomain::Val; kwargs...) - return $f(MatricizeStyle(A), A, ndims_codomain; kwargs...) - end end end # `svd_trunc` matches the three-output SVD but additionally surfaces the truncation error # `ϵ` (the 2-norm of the discarded singular values, computed by MatrixAlgebraKit without # catastrophic cancellation), so it is spelled out here rather than sharing the loop above. -function svd_trunc(style::MatricizeStyle, A, ndims_codomain::Val; kwargs...) - A_mat = matricize(style, A, ndims_codomain) - U, S, Vᴴ, ϵ = MatrixAlgebraKit.svd_trunc(A_mat; kwargs...) - axes_codomain, axes_domain = bipartition_axes(axes(A), ndims_codomain) +function matricized( + ::typeof(svd_trunc), style::MatricizeStyle, A_mat, owned::Bool, + axes_codomain, axes_domain; kwargs... + ) + U, S, Vᴴ, ϵ = maybe_donate( + MatrixAlgebraKit.svd_trunc, MatrixAlgebraKit.svd_trunc!, A_mat, owned; kwargs... + ) return unmatricize(style, U, axes_codomain, (conj(axes(U, ndims(U))),)), S, unmatricize(style, Vᴴ, (axes(Vᴴ, 1),), axes_domain), ϵ end -function svd_trunc(A, ndims_codomain::Val; kwargs...) - return svd_trunc(MatricizeStyle(A), A, ndims_codomain; kwargs...) -end # Eigendecomposition: `D` is the rank × rank spectrum and `V` carries the codomain axes plus a # trailing rank axis. `D` is returned bare (its axis is the internal bond, so there is nothing to # unfold); `V` is unmatricized back to the array type, as in `svd_*`. for f in (:eigh_full, :eig_full, :eigh_trunc, :eig_trunc) @eval begin - function $f(style::MatricizeStyle, A, ndims_codomain::Val; kwargs...) - A_mat = matricize(style, A, ndims_codomain) - D, V = MatrixAlgebraKit.$f(A_mat; kwargs...) - axes_codomain = first(bipartition(axes(A), ndims_codomain)) + function matricized( + ::typeof($f), style::MatricizeStyle, A_mat, owned::Bool, + axes_codomain, axes_domain; kwargs... + ) + D, V = maybe_donate( + MatrixAlgebraKit.$f, MatrixAlgebraKit.$(Symbol(f, :!)), A_mat, owned; + kwargs... + ) return D, unmatricize(style, V, axes_codomain, (conj(axes(V, ndims(V))),)) end - function $f(A, ndims_codomain::Val; kwargs...) - return $f(MatricizeStyle(A), A, ndims_codomain; kwargs...) - end end end # Spectrum-only factorizations returning a vector of singular values / eigenvalues. for f in (:svd_vals, :eigh_vals, :eig_vals) @eval begin - function $f(style::MatricizeStyle, A, ndims_codomain::Val; kwargs...) - A_mat = matricize(style, A, ndims_codomain) + function matricized( + ::typeof($f), style::MatricizeStyle, A_mat, owned::Bool, + axes_codomain, axes_domain; kwargs... + ) return MatrixAlgebraKit.$f(A_mat; kwargs...) end - function $f(A, ndims_codomain::Val; kwargs...) - return $f(MatricizeStyle(A), A, ndims_codomain; kwargs...) - end end end @@ -500,11 +541,14 @@ function left_null!!(A, ndims_codomain::Val; kwargs...) return left_null!!(MatricizeStyle(A), A, ndims_codomain; kwargs...) end -function left_null(style::MatricizeStyle, A, ndims_codomain::Val; kwargs...) - return left_null!!(style, copy(A), ndims_codomain; kwargs...) -end -function left_null(A, ndims_codomain::Val; kwargs...) - return left_null!!(copy(A), ndims_codomain; kwargs...) +function matricized( + ::typeof(left_null), style::MatricizeStyle, A_mat, owned::Bool, + axes_codomain, axes_domain; kwargs... + ) + N = maybe_donate( + MatrixAlgebraKit.left_null, MatrixAlgebraKit.left_null!, A_mat, owned; kwargs... + ) + return unmatricize(style, N, axes_codomain, (conj(axes(N, ndims(N))),)) end """ @@ -537,11 +581,15 @@ function right_null!!(A, ndims_codomain::Val; kwargs...) return right_null!!(MatricizeStyle(A), A, ndims_codomain; kwargs...) end -function right_null(style::MatricizeStyle, A, ndims_codomain::Val; kwargs...) - return right_null!!(style, copy(A), ndims_codomain; kwargs...) -end -function right_null(A, ndims_codomain::Val; kwargs...) - return right_null!!(copy(A), ndims_codomain; kwargs...) +function matricized( + ::typeof(right_null), style::MatricizeStyle, A_mat, owned::Bool, + axes_codomain, axes_domain; kwargs... + ) + Nᴴ = maybe_donate( + MatrixAlgebraKit.right_null, MatrixAlgebraKit.right_null!, A_mat, owned; + kwargs... + ) + return unmatricize(style, Nᴴ, (axes(Nᴴ, 1),), axes_domain) end """ @@ -593,13 +641,14 @@ function gram_eigh_full!!(A, ndims_codomain::Val; kwargs...) return gram_eigh_full!!(MatricizeStyle(A), A, ndims_codomain; kwargs...) end -function gram_eigh_full( - style::MatricizeStyle, A, ndims_codomain::Val; kwargs... +# The non-mutating matrix-level `gram_eigh_full` copies its input internally (through +# `MatrixAlgebraKit.eigh_full`), so the maybe-alias `A_mat` is consumed read-only. +function matricized( + ::typeof(gram_eigh_full), style::MatricizeStyle, A_mat, owned::Bool, + axes_codomain, axes_domain; kwargs... ) - return gram_eigh_full!!(style, copy(A), ndims_codomain; kwargs...) -end -function gram_eigh_full(A, ndims_codomain::Val; kwargs...) - return gram_eigh_full!!(copy(A), ndims_codomain; kwargs...) + X = MatrixAlgebra.gram_eigh_full(A_mat; kwargs...) + return unmatricize(style, X, axes_codomain, (conj(axes(X, ndims(X))),)) end """ @@ -655,13 +704,13 @@ function gram_eigh_full_with_pinv!!(A, ndims_codomain::Val; kwargs...) return gram_eigh_full_with_pinv!!(MatricizeStyle(A), A, ndims_codomain; kwargs...) end -function gram_eigh_full_with_pinv( - style::MatricizeStyle, A, ndims_codomain::Val; kwargs... +function matricized( + ::typeof(gram_eigh_full_with_pinv), style::MatricizeStyle, A_mat, owned::Bool, + axes_codomain, axes_domain; kwargs... ) - return gram_eigh_full_with_pinv!!(style, copy(A), ndims_codomain; kwargs...) -end -function gram_eigh_full_with_pinv(A, ndims_codomain::Val; kwargs...) - return gram_eigh_full_with_pinv!!(copy(A), ndims_codomain; kwargs...) + X, Y = MatrixAlgebra.gram_eigh_full_with_pinv(A_mat; kwargs...) + return unmatricize(style, X, axes_codomain, (conj(axes(X, ndims(X))),)), + unmatricize(style, Y, (axes(Y, 1),), axes_codomain) end """ @@ -712,15 +761,13 @@ invsqrth_safe for f in (:sqrth_safe, :invsqrth_safe) @eval begin - function $f(style::MatricizeStyle, A, ndims_codomain::Val; kwargs...) - A_mat = matricize(style, A, ndims_codomain) + function matricized( + ::typeof($f), style::MatricizeStyle, A_mat, owned::Bool, + axes_codomain, axes_domain; kwargs... + ) P_mat = MatrixAlgebra.$f(A_mat; kwargs...) - axes_codomain, axes_domain = bipartition_axes(axes(A), ndims_codomain) return unmatricize(style, P_mat, axes_codomain, axes_domain) end - function $f(A, ndims_codomain::Val; kwargs...) - return $f(MatricizeStyle(A), A, ndims_codomain; kwargs...) - end end end @@ -737,15 +784,13 @@ See also `MatrixAlgebraKit.project_hermitian`. """ project_hermitian -function project_hermitian(style::MatricizeStyle, A, ndims_codomain::Val; kwargs...) - A_mat = matricize(style, A, ndims_codomain) +function matricized( + ::typeof(project_hermitian), style::MatricizeStyle, A_mat, owned::Bool, + axes_codomain, axes_domain; kwargs... + ) H_mat = MatrixAlgebraKit.project_hermitian(A_mat; kwargs...) - axes_codomain, axes_domain = bipartition_axes(axes(A), ndims_codomain) return unmatricize(style, H_mat, axes_codomain, axes_domain) end -function project_hermitian(A, ndims_codomain::Val; kwargs...) - return project_hermitian(MatricizeStyle(A), A, ndims_codomain; kwargs...) -end """ sqrth_invsqrth_safe(A, labels_A, labels_codomain, labels_domain; kwargs...) -> P, Pinv @@ -767,16 +812,14 @@ See also [`MatrixAlgebra.sqrth_invsqrth_safe`](@ref). """ sqrth_invsqrth_safe -function sqrth_invsqrth_safe(style::MatricizeStyle, A, ndims_codomain::Val; kwargs...) - A_mat = matricize(style, A, ndims_codomain) +function matricized( + ::typeof(sqrth_invsqrth_safe), style::MatricizeStyle, A_mat, owned::Bool, + axes_codomain, axes_domain; kwargs... + ) P_mat, Pinv_mat = MatrixAlgebra.sqrth_invsqrth_safe(A_mat; kwargs...) - axes_codomain, axes_domain = bipartition_axes(axes(A), ndims_codomain) return unmatricize(style, P_mat, axes_codomain, axes_domain), unmatricize(style, Pinv_mat, axes_codomain, axes_domain) end -function sqrth_invsqrth_safe(A, ndims_codomain::Val; kwargs...) - return sqrth_invsqrth_safe(MatricizeStyle(A), A, ndims_codomain; kwargs...) -end """ TensorAlgebra.one(A, labels_A, labels_codomain, labels_domain) -> Id @@ -840,3 +883,33 @@ end function one(A, ndims_codomain::Val; kwargs...) return one!!(copy(A), ndims_codomain; kwargs...) end + +# `one` stays off the shared `matricized` wrappers: `one!!` is its own overload point (a +# `TensorMap` backend fills the identity through TensorKit rather than MatrixAlgebraKit), and +# the prototype only supplies shape, so the permuted forms hand `one!!` the `bipermutedims` +# result directly — always a fresh buffer, by the `bipermutedims` copy convention. +function one( + style::MatricizeStyle, A, + perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; + kwargs... + ) + A_perm = bipermutedims(A, perm_codomain, perm_domain) + return one!!(style, A_perm, Val(length(perm_codomain)); kwargs...) +end +function one( + A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs... + ) + return one(MatricizeStyle(A), A, perm_codomain, perm_domain; kwargs...) +end +function one( + style::MatricizeStyle, A, labels_A, labels_codomain, labels_domain; kwargs... + ) + perm_codomain, perm_domain = + biperm(Tuple.((labels_A, labels_codomain, labels_domain))...) + return one(style, A, perm_codomain, perm_domain; kwargs...) +end +function one(A, labels_A, labels_codomain, labels_domain; kwargs...) + perm_codomain, perm_domain = + biperm(Tuple.((labels_A, labels_codomain, labels_domain))...) + return one(A, perm_codomain, perm_domain; kwargs...) +end diff --git a/src/matricize.jl b/src/matricize.jl index ed6cd1c0..ffaabc96 100644 --- a/src/matricize.jl +++ b/src/matricize.jl @@ -68,9 +68,12 @@ function bipermutedims!( end # ===================================== matricize ======================================== -# TBD settle copy/not copy convention -# matrix factorizations assume copy -# maybe: copy=false kwarg +# Copy convention: `bipermutedims`/`permutedims` always copy (Base `permutedims` semantics). +# `matricize`/`matricizeperm`/`matricizeopperm` are the maybe-alias tier — the result may be a +# view of the input or a fresh gather, and callers must treat it as read-only. A consumer that +# mutates takes an explicit copy (`MatrixAlgebraKit.copy_input` or `copy`); the factorization +# wrappers donate a provably owned matricization to the mutating matrix-level entries (see +# `maybe_donate` in `factorizations.jl`). # This is the primary function that should be overloaded for new matricize styles. # This assumes the permutation was already performed. @@ -132,6 +135,23 @@ function matricizeperm( return matricizeperm(style, a, to_permblocks(a, (perm_codomain, perm_domain))...) end +# `matricizeperm` for the factorization and matrix-function wrappers: same maybe-alias +# contract, but a codomain/domain swap takes the permuted copy instead of the lazy `transpose` +# view — matrix-level backends (LAPACK through MatrixAlgebraKit) require the matricized layout +# itself, not just any matrix-shaped view. +function matricize_input( + style::MatricizeStyle, a, + perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}} + ) + ndims(a) == length(perm_codomain) + length(perm_domain) || + throw(ArgumentError("Invalid bipermutation")) + kind = matricizekind(style, perm_codomain, perm_domain) + kind == ReshapeMatricizeKind && + return matricize(style, a, Val(length(perm_codomain))) + a_perm = bipermutedims(a, perm_codomain, perm_domain) + return matricize(style, a_perm, Val(length(perm_codomain))) +end + # ================================== matricizeopperm ===================================== """ diff --git a/src/matrixfunctions.jl b/src/matrixfunctions.jl index 5ccdc460..3d571063 100644 --- a/src/matrixfunctions.jl +++ b/src/matrixfunctions.jl @@ -31,13 +31,18 @@ const MATRIX_FUNCTIONS = [ :acoth, ] +# The wrappers share the factorization machinery: `matricized` cores below consume the +# maybe-alias matricization read-only (the matrix functions allocate their own outputs), so +# the permuted forms skip the eager `bipermutedims` copy (see `factorizations.jl`). for f in MATRIX_FUNCTIONS @eval begin function $f(style::MatricizeStyle, a, ndims_codomain::Val; kwargs...) a_mat = matricize(style, a, ndims_codomain) - fa_mat = Base.$f(a_mat; kwargs...) axes_codomain, axes_domain = bipartition_axes(axes(a), ndims_codomain) - return unmatricize(style, fa_mat, axes_codomain, axes_domain) + return matricized( + $f, style, a_mat, isdetached(a_mat, a), + axes_codomain, axes_domain; kwargs... + ) end function $f(a, ndims_codomain::Val; kwargs...) return $f(MatricizeStyle(a), a, ndims_codomain; kwargs...) @@ -48,16 +53,30 @@ for f in MATRIX_FUNCTIONS perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs... ) - a_perm = bipermutedims(a, perm_codomain, perm_domain) - return $f(style, a_perm, Val(length(perm_codomain)); kwargs...) + a_mat = matricize_input(style, a, perm_codomain, perm_domain) + axes_codomain, axes_domain = bipartition_axes( + map(i -> axes(a, i), (perm_codomain..., perm_domain...)), + Val(length(perm_codomain)) + ) + return matricized( + $f, style, a_mat, isdetached(a_mat, a), + axes_codomain, axes_domain; kwargs... + ) end function $f( a, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs... ) - a_perm = bipermutedims(a, perm_codomain, perm_domain) - return $f(a_perm, Val(length(perm_codomain)); kwargs...) + return $f(MatricizeStyle(a), a, perm_codomain, perm_domain; kwargs...) + end + + function matricized( + ::typeof($f), style::MatricizeStyle, a_mat, owned::Bool, + axes_codomain, axes_domain; kwargs... + ) + fa_mat = Base.$f(a_mat; kwargs...) + return unmatricize(style, fa_mat, axes_codomain, axes_domain) end function $f( diff --git a/test/test_factorizations.jl b/test/test_factorizations.jl index a933ee0a..45357206 100644 --- a/test/test_factorizations.jl +++ b/test/test_factorizations.jl @@ -418,3 +418,45 @@ end @test TensorAlgebra.tr(A, (1, 2), (3, 4)) ≈ LinearAlgebra.tr(m) @test TensorAlgebra.tr(A, Val(2)) ≈ LinearAlgebra.tr(m) end + +# Permuted entry points: matricization and buffer donation +# -------------------------------------------------------- +# The permuted forms matricize directly (`matricize_input`, no eager `bipermutedims` copy) +# and donate an owned matricization to the mutating MatrixAlgebraKit entries, so pin them +# to the reference permute-then-`Val` path and check the caller's array is never mutated, +# through the identity, permuted-copy, and codomain/domain-swap matricizations alike. +@testset "Permuted forms match permute-then-matricize ($T)" for T in elts + A = randn(T, 2, 3, 4) + Acopy = copy(A) + for (perm_codomain, perm_domain) in + (((1, 2), (3,)), ((3, 1), (2,)), ((3,), (1, 2)), ((2,), (3, 1))) + A_perm = TensorAlgebra.bipermutedims(A, perm_codomain, perm_domain) + for f in ( + qr_compact, lq_compact, left_orth, right_orth, + svd_compact, svd_trunc, svd_vals, left_null, right_null, + ) + F = f(A, perm_codomain, perm_domain) + F_ref = f(A_perm, Val(length(perm_codomain))) + Fs = F isa Tuple ? F : (F,) + F_refs = F_ref isa Tuple ? F_ref : (F_ref,) + @test all(map(==, Fs, F_refs)) + @test A == Acopy + end + end + B = randn(T, 2, 3, 2, 3) + Bcopy = copy(B) + for (perm_codomain, perm_domain) in + (((1, 2), (3, 4)), ((3, 4), (1, 2)), ((2, 3), (4, 1))) + B_perm = TensorAlgebra.bipermutedims(B, perm_codomain, perm_domain) + for f in (eig_full, eig_vals) + F = f(B, perm_codomain, perm_domain) + F_ref = f(B_perm, Val(2)) + Fs = F isa Tuple ? F : (F,) + F_refs = F_ref isa Tuple ? F_ref : (F_ref,) + @test all(map(==, Fs, F_refs)) + @test B == Bcopy + end + @test TensorAlgebra.tr(B, perm_codomain, perm_domain) ≈ + TensorAlgebra.tr(B_perm, Val(2)) + end +end diff --git a/test/test_matrixfunctions.jl b/test/test_matrixfunctions.jl index 23ba84ae..4ff9f06d 100644 --- a/test/test_matrixfunctions.jl +++ b/test/test_matrixfunctions.jl @@ -20,6 +20,15 @@ using Test: @test, @testset fa = TensorAlgebra.$f(a, Val(2)) fa′ = reshape($f(reshape(a, (4, 4))), (2, 2, 2, 2)) @test fa ≈ fa′ + # Codomain/domain swap and identity bipermutations: the storage-order fast + # paths of the wrapper's matricization, which must not touch `a`. + acopy = copy(a) + fa_swap = TensorAlgebra.$f(a, (3, 4), (1, 2)) + fa_swap′ = + reshape($f(reshape(permutedims(a, (3, 4, 1, 2)), (4, 4))), (2, 2, 2, 2)) + @test fa_swap ≈ fa_swap′ + @test TensorAlgebra.$f(a, (1, 2), (3, 4)) ≈ fa + @test a == acopy end end end From a8f863ed6658cef6cc37d44cc988e54ee80aa1c8 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Sun, 23 Aug 2026 09:32:09 -0400 Subject: [PATCH 2/5] Replace matricization aliasing inference with declared sharing --- src/contract/contract_matricize.jl | 32 ++--- src/factorizations.jl | 210 ++++++++++++++++------------- src/matricize.jl | 116 ++++++---------- src/matrixfunctions.jl | 34 ++--- test/test_factorizations.jl | 112 +++++++++++---- test/test_matricize.jl | 109 +++++++-------- test/test_matrixalgebra.jl | 2 +- 7 files changed, 316 insertions(+), 299 deletions(-) diff --git a/src/contract/contract_matricize.jl b/src/contract/contract_matricize.jl index 3ade05ab..a1ff5a5f 100644 --- a/src/contract/contract_matricize.jl +++ b/src/contract/contract_matricize.jl @@ -23,32 +23,24 @@ function contractopadd!( algorithm.right_matricize_style, op2, a2, biperm2_codomain, biperm2_domain ) output_style = algorithm.output_matricize_style - if iszero(β) && !matricizepermaliases(output_style, invperm_codomain, invperm_domain) - # `β` is a strong zero and matricizing `a_dest` would only build a detached copy that - # `mul!` immediately overwrites, so skip that gather: let the matmul allocate its matrix - # result directly and scatter it into `a_dest`. Every coupled-sector block is - # materialized (the matmul zeros the ones it does not reach), so the scatter overwrites - # `a_dest` in full. + a_dest_mat = trymatricizeview(output_style, a_dest, invperm_codomain, invperm_domain) + if !isnothing(a_dest_mat) + # The matricization shares `a_dest`'s memory, so the matmul is the whole operation. + mul!(a_dest_mat, a1_mat, a2_mat, α, β) + elseif iszero(β) + # `β` is a strong zero, so `a_dest`'s current data is irrelevant: let the matmul + # allocate its matrix result and scatter it into `a_dest`. Every coupled-sector block + # is materialized (the matmul zeros the ones it does not reach), so the scatter + # overwrites `a_dest` in full. a_dest_mat = a1_mat * a2_mat isone(α) || scale!(a_dest_mat, α) unmatricizeperm!(output_style, a_dest, a_dest_mat, invperm_codomain, invperm_domain) else - # Matricize the destination and multiply straight into it: a no-op for an aligned or - # transposed dense output (a view aliasing `a_dest`, so `mul!` writes through and we - # are done), a fresh permuted copy otherwise. Either way `matricize` seeds `a_dest_mat` - # with `a_dest`'s current contents, so `β` rides on the `mul!` and a detached copy is - # written back with a plain overwrite. + # `a_dest`'s data contributes through `β`, so gather it, multiply into the gathered + # copy, and scatter back. a_dest_mat = matricizeperm(output_style, a_dest, invperm_codomain, invperm_domain) mul!(a_dest_mat, a1_mat, a2_mat, α, β) - if !Base.mightalias(a_dest_mat, a_dest) - unmatricizeperm!( - output_style, - a_dest, - a_dest_mat, - invperm_codomain, - invperm_domain - ) - end + unmatricizeperm!(output_style, a_dest, a_dest_mat, invperm_codomain, invperm_domain) end return a_dest end diff --git a/src/factorizations.jl b/src/factorizations.jl index 01c2383f..7bdf53f2 100644 --- a/src/factorizations.jl +++ b/src/factorizations.jl @@ -9,64 +9,101 @@ using MatrixAlgebraKit: MatrixAlgebraKit # bond is dualized to codomain-facing form (`conj`, a no-op on a dense axis) when it lands on the # domain side of the reconstruction, matching the `unmatricize`/`similar_map` axis convention. -# Whether the matricized form `A_mat` is detached from `A`: fresh storage the wrapper owns, so -# a mutating consumer cannot corrupt the caller's data. Aliasing is only tracked (through -# `Base.dataids`) for `AbstractArray`s; any other tensor type (such as a `TensorMap`, whose -# trivial `permute` returns the input itself) is conservatively treated as attached. -isdetached(A_mat::AbstractArray, A::AbstractArray) = !Base.mightalias(A_mat, A) -isdetached(A_mat, A) = false - -# Skip the copy inside the non-mutating MatrixAlgebraKit entry (`f(A) = f!(copy_input(f, A))`) -# when the wrapper owns `A_mat` and it already has the floating-point eltype `copy_input` -# would produce. When `A_mat` may alias the caller's input (a dense reshape or a stored graded -# matrix), the copying entry is required: mutating `A_mat` would corrupt that input. -function maybe_donate(f, f!, A_mat, owned::Bool; kwargs...) - owned && eltype(A_mat) === float(eltype(A_mat)) && return f!(A_mat; kwargs...) - return f(A_mat; kwargs...) -end - -# `matricized(f, style, A_mat, owned, axes_codomain, axes_domain; kwargs...)` is the shared -# matrix-level body of the wrapper `f`: apply the matrix-level `f` to the matricized input +# `unmatricize_factors(f, style, A_mat, axes_codomain, axes_domain; kwargs...)` is the shared +# matrix-level body of the wrapper `f`: apply the matrix-level function to the matricized input # `A_mat` and unfold the outputs with the bipartitioned axes (in the `unmatricize` convention, -# domain axes un-dualized). `owned` (see `isdetached`) gates donating `A_mat` to the mutating -# matrix-level entry. The wrappers produce `A_mat` with the maybe-alias `matricize`/ -# `matricize_input` spelling, so the permuted forms skip the eager `bipermutedims` copy. +# domain axes un-dualized). +# +# Owned tier: the matrix-level entries mutate their input, so the perm form materializes an +# owned matricization following MatrixAlgebraKit's `f(A) = f!(copy_input(f, A))` convention — +# a memory-sharing matricization is materialized through `MatrixAlgebraKit.copy_input`, while +# the gathered bipermutation is fresh by construction and is donated directly (with +# `copy_input` still applied when the eltype must change) — and the cores call the mutating +# entry unconditionally. for f in ( :qr_compact, :qr_full, :lq_compact, :lq_full, :left_polar, :right_polar, :left_orth, :right_orth, :svd_compact, :svd_full, :svd_trunc, :svd_vals, :eigh_full, :eig_full, :eigh_trunc, :eig_trunc, :eigh_vals, :eig_vals, - :left_null, :right_null, :gram_eigh_full, :gram_eigh_full_with_pinv, - :sqrth_safe, :invsqrth_safe, :sqrth_invsqrth_safe, :project_hermitian, + :left_null, :right_null, :project_hermitian, ) @eval begin - function $f(style::MatricizeStyle, A, ndims_codomain::Val; kwargs...) - A_mat = matricize(style, A, ndims_codomain) - axes_codomain, axes_domain = bipartition_axes(axes(A), ndims_codomain) - return matricized( - $f, style, A_mat, isdetached(A_mat, A), - axes_codomain, axes_domain; kwargs... + function $f( + style::MatricizeStyle, A, + perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; + kwargs... + ) + ndims(A) == length(perm_codomain) + length(perm_domain) || + throw(ArgumentError("Invalid bipermutation")) + A_shared = trymatricizeview(style, A, perm_codomain, perm_domain) + A_mat = if !isnothing(A_shared) + MatrixAlgebraKit.copy_input(MatrixAlgebraKit.$f, A_shared) + else + A_perm = bipermutedims(A, perm_codomain, perm_domain) + A_gather = matricize(style, A_perm, Val(length(perm_codomain))) + if eltype(A_gather) === float(eltype(A_gather)) + A_gather + else + MatrixAlgebraKit.copy_input(MatrixAlgebraKit.$f, A_gather) + end + end + axes_codomain, axes_domain = bipartition_axes( + map(i -> axes(A, i), (perm_codomain..., perm_domain...)), + Val(length(perm_codomain)) + ) + return unmatricize_factors( + $f, style, A_mat, axes_codomain, axes_domain; kwargs... ) end - function $f(A, ndims_codomain::Val; kwargs...) - return $f(MatricizeStyle(A), A, ndims_codomain; kwargs...) - end + end +end +# Read-only tier: the matrix-level entries never mutate their input (they copy internally), so +# the perm form consumes the maybe-alias `matricizeperm` matricization directly. +for f in ( + :gram_eigh_full, :gram_eigh_full_with_pinv, + :sqrth_safe, :invsqrth_safe, :sqrth_invsqrth_safe, + ) + @eval begin function $f( style::MatricizeStyle, A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs... ) - A_mat = matricize_input(style, A, perm_codomain, perm_domain) + A_mat = matricizeperm(style, A, perm_codomain, perm_domain) axes_codomain, axes_domain = bipartition_axes( map(i -> axes(A, i), (perm_codomain..., perm_domain...)), Val(length(perm_codomain)) ) - return matricized( - $f, style, A_mat, isdetached(A_mat, A), - axes_codomain, axes_domain; kwargs... + return unmatricize_factors( + $f, style, A_mat, axes_codomain, axes_domain; kwargs... ) end + end +end + +# The `Val`, style-inferring, and labels forms of both tiers are thin forwarders into the perm +# form, at the identity bipermutation for the `Val` form. +for f in ( + :qr_compact, :qr_full, :lq_compact, :lq_full, + :left_polar, :right_polar, :left_orth, :right_orth, + :svd_compact, :svd_full, :svd_trunc, :svd_vals, + :eigh_full, :eig_full, :eigh_trunc, :eig_trunc, :eigh_vals, :eig_vals, + :left_null, :right_null, :gram_eigh_full, :gram_eigh_full_with_pinv, + :sqrth_safe, :invsqrth_safe, :sqrth_invsqrth_safe, :project_hermitian, + ) + @eval begin + function $f(style::MatricizeStyle, A, ndims_codomain::Val{K}; kwargs...) where {K} + return $f( + style, A, + ntuple(identity, ndims_codomain), + ntuple(i -> K + i, Val(ndims(A) - K)); + kwargs... + ) + end + function $f(A, ndims_codomain::Val; kwargs...) + return $f(MatricizeStyle(A), A, ndims_codomain; kwargs...) + end function $f( A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; @@ -74,7 +111,6 @@ for f in ( ) return $f(MatricizeStyle(A), A, perm_codomain, perm_domain; kwargs...) end - function $f( style::MatricizeStyle, A, labels_A, labels_codomain, labels_domain; kwargs... @@ -98,14 +134,11 @@ for f in ( :left_polar, :right_polar, :left_orth, :right_orth, ) @eval begin - function matricized( - ::typeof($f), style::MatricizeStyle, A_mat, owned::Bool, + function unmatricize_factors( + ::typeof($f), style::MatricizeStyle, A_mat, axes_codomain, axes_domain; kwargs... ) - X, Y = maybe_donate( - MatrixAlgebraKit.$f, MatrixAlgebraKit.$(Symbol(f, :!)), A_mat, owned; - kwargs... - ) + X, Y = MatrixAlgebraKit.$(Symbol(f, :!))(A_mat; kwargs...) return unmatricize(style, X, axes_codomain, (conj(axes(X, ndims(X))),)), unmatricize(style, Y, (axes(Y, 1),), axes_domain) end @@ -297,14 +330,11 @@ right_orth # rank × rank spectrum, and `Vᴴ` carries a leading rank axis plus the domain axes. for f in (:svd_compact, :svd_full) @eval begin - function matricized( - ::typeof($f), style::MatricizeStyle, A_mat, owned::Bool, + function unmatricize_factors( + ::typeof($f), style::MatricizeStyle, A_mat, axes_codomain, axes_domain; kwargs... ) - U, S, Vᴴ = maybe_donate( - MatrixAlgebraKit.$f, MatrixAlgebraKit.$(Symbol(f, :!)), A_mat, owned; - kwargs... - ) + U, S, Vᴴ = MatrixAlgebraKit.$(Symbol(f, :!))(A_mat; kwargs...) return unmatricize(style, U, axes_codomain, (conj(axes(U, ndims(U))),)), S, unmatricize(style, Vᴴ, (axes(Vᴴ, 1),), axes_domain) @@ -315,13 +345,11 @@ end # `svd_trunc` matches the three-output SVD but additionally surfaces the truncation error # `ϵ` (the 2-norm of the discarded singular values, computed by MatrixAlgebraKit without # catastrophic cancellation), so it is spelled out here rather than sharing the loop above. -function matricized( - ::typeof(svd_trunc), style::MatricizeStyle, A_mat, owned::Bool, +function unmatricize_factors( + ::typeof(svd_trunc), style::MatricizeStyle, A_mat, axes_codomain, axes_domain; kwargs... ) - U, S, Vᴴ, ϵ = maybe_donate( - MatrixAlgebraKit.svd_trunc, MatrixAlgebraKit.svd_trunc!, A_mat, owned; kwargs... - ) + U, S, Vᴴ, ϵ = MatrixAlgebraKit.svd_trunc!(A_mat; kwargs...) return unmatricize(style, U, axes_codomain, (conj(axes(U, ndims(U))),)), S, unmatricize(style, Vᴴ, (axes(Vᴴ, 1),), axes_domain), @@ -333,14 +361,11 @@ end # unfold); `V` is unmatricized back to the array type, as in `svd_*`. for f in (:eigh_full, :eig_full, :eigh_trunc, :eig_trunc) @eval begin - function matricized( - ::typeof($f), style::MatricizeStyle, A_mat, owned::Bool, + function unmatricize_factors( + ::typeof($f), style::MatricizeStyle, A_mat, axes_codomain, axes_domain; kwargs... ) - D, V = maybe_donate( - MatrixAlgebraKit.$f, MatrixAlgebraKit.$(Symbol(f, :!)), A_mat, owned; - kwargs... - ) + D, V = MatrixAlgebraKit.$(Symbol(f, :!))(A_mat; kwargs...) return D, unmatricize(style, V, axes_codomain, (conj(axes(V, ndims(V))),)) end @@ -350,11 +375,11 @@ end # Spectrum-only factorizations returning a vector of singular values / eigenvalues. for f in (:svd_vals, :eigh_vals, :eig_vals) @eval begin - function matricized( - ::typeof($f), style::MatricizeStyle, A_mat, owned::Bool, + function unmatricize_factors( + ::typeof($f), style::MatricizeStyle, A_mat, axes_codomain, axes_domain; kwargs... ) - return MatrixAlgebraKit.$f(A_mat; kwargs...) + return MatrixAlgebraKit.$(Symbol(f, :!))(A_mat; kwargs...) end end end @@ -541,13 +566,11 @@ function left_null!!(A, ndims_codomain::Val; kwargs...) return left_null!!(MatricizeStyle(A), A, ndims_codomain; kwargs...) end -function matricized( - ::typeof(left_null), style::MatricizeStyle, A_mat, owned::Bool, +function unmatricize_factors( + ::typeof(left_null), style::MatricizeStyle, A_mat, axes_codomain, axes_domain; kwargs... ) - N = maybe_donate( - MatrixAlgebraKit.left_null, MatrixAlgebraKit.left_null!, A_mat, owned; kwargs... - ) + N = MatrixAlgebraKit.left_null!(A_mat; kwargs...) return unmatricize(style, N, axes_codomain, (conj(axes(N, ndims(N))),)) end @@ -581,14 +604,11 @@ function right_null!!(A, ndims_codomain::Val; kwargs...) return right_null!!(MatricizeStyle(A), A, ndims_codomain; kwargs...) end -function matricized( - ::typeof(right_null), style::MatricizeStyle, A_mat, owned::Bool, +function unmatricize_factors( + ::typeof(right_null), style::MatricizeStyle, A_mat, axes_codomain, axes_domain; kwargs... ) - Nᴴ = maybe_donate( - MatrixAlgebraKit.right_null, MatrixAlgebraKit.right_null!, A_mat, owned; - kwargs... - ) + Nᴴ = MatrixAlgebraKit.right_null!(A_mat; kwargs...) return unmatricize(style, Nᴴ, (axes(Nᴴ, 1),), axes_domain) end @@ -641,10 +661,8 @@ function gram_eigh_full!!(A, ndims_codomain::Val; kwargs...) return gram_eigh_full!!(MatricizeStyle(A), A, ndims_codomain; kwargs...) end -# The non-mutating matrix-level `gram_eigh_full` copies its input internally (through -# `MatrixAlgebraKit.eigh_full`), so the maybe-alias `A_mat` is consumed read-only. -function matricized( - ::typeof(gram_eigh_full), style::MatricizeStyle, A_mat, owned::Bool, +function unmatricize_factors( + ::typeof(gram_eigh_full), style::MatricizeStyle, A_mat, axes_codomain, axes_domain; kwargs... ) X = MatrixAlgebra.gram_eigh_full(A_mat; kwargs...) @@ -704,8 +722,8 @@ function gram_eigh_full_with_pinv!!(A, ndims_codomain::Val; kwargs...) return gram_eigh_full_with_pinv!!(MatricizeStyle(A), A, ndims_codomain; kwargs...) end -function matricized( - ::typeof(gram_eigh_full_with_pinv), style::MatricizeStyle, A_mat, owned::Bool, +function unmatricize_factors( + ::typeof(gram_eigh_full_with_pinv), style::MatricizeStyle, A_mat, axes_codomain, axes_domain; kwargs... ) X, Y = MatrixAlgebra.gram_eigh_full_with_pinv(A_mat; kwargs...) @@ -761,8 +779,8 @@ invsqrth_safe for f in (:sqrth_safe, :invsqrth_safe) @eval begin - function matricized( - ::typeof($f), style::MatricizeStyle, A_mat, owned::Bool, + function unmatricize_factors( + ::typeof($f), style::MatricizeStyle, A_mat, axes_codomain, axes_domain; kwargs... ) P_mat = MatrixAlgebra.$f(A_mat; kwargs...) @@ -784,11 +802,11 @@ See also `MatrixAlgebraKit.project_hermitian`. """ project_hermitian -function matricized( - ::typeof(project_hermitian), style::MatricizeStyle, A_mat, owned::Bool, +function unmatricize_factors( + ::typeof(project_hermitian), style::MatricizeStyle, A_mat, axes_codomain, axes_domain; kwargs... ) - H_mat = MatrixAlgebraKit.project_hermitian(A_mat; kwargs...) + H_mat = MatrixAlgebraKit.project_hermitian!(A_mat; kwargs...) return unmatricize(style, H_mat, axes_codomain, axes_domain) end @@ -812,8 +830,8 @@ See also [`MatrixAlgebra.sqrth_invsqrth_safe`](@ref). """ sqrth_invsqrth_safe -function matricized( - ::typeof(sqrth_invsqrth_safe), style::MatricizeStyle, A_mat, owned::Bool, +function unmatricize_factors( + ::typeof(sqrth_invsqrth_safe), style::MatricizeStyle, A_mat, axes_codomain, axes_domain; kwargs... ) P_mat, Pinv_mat = MatrixAlgebra.sqrth_invsqrth_safe(A_mat; kwargs...) @@ -864,13 +882,17 @@ function one!!(A, ndims_codomain::Val; kwargs...) return one!!(MatricizeStyle(A), A, ndims_codomain; kwargs...) end -# In-place identity fill: writes the identity into `A` and returns it. Matricizes `A`, fills the -# fused matrix with the identity, and — when the matricized form is a detached copy (a graded -# gather) rather than a view aliasing `A` (a dense reshape) — scatters it back with `unmatricize!`. +# In-place identity fill: writes the identity into `A` and returns it. Fills a memory-sharing +# matricization directly when the style declares one at this split, and otherwise fills a +# gathered matrix and scatters it back with `unmatricize!`. function one!(style::MatricizeStyle, A, ndims_codomain::Val; kwargs...) + A_mat = trymatricizeview(style, A, ndims_codomain) + if !isnothing(A_mat) + MatrixAlgebraKit.one!(A_mat) + return A + end A_mat = matricize(style, A, ndims_codomain) MatrixAlgebraKit.one!(A_mat) - Base.mightalias(A_mat, A) && return A return unmatricize!(A, A_mat, ndims_codomain) end function one!(A, ndims_codomain::Val; kwargs...) @@ -884,10 +906,8 @@ function one(A, ndims_codomain::Val; kwargs...) return one!!(copy(A), ndims_codomain; kwargs...) end -# `one` stays off the shared `matricized` wrappers: `one!!` is its own overload point (a -# `TensorMap` backend fills the identity through TensorKit rather than MatrixAlgebraKit), and -# the prototype only supplies shape, so the permuted forms hand `one!!` the `bipermutedims` -# result directly — always a fresh buffer, by the `bipermutedims` copy convention. +# `one` stays off the shared factorization wrappers: `one!!` is its own overload point (a +# `TensorMap` backend fills the identity through TensorKit rather than MatrixAlgebraKit). function one( style::MatricizeStyle, A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; diff --git a/src/matricize.jl b/src/matricize.jl index ffaabc96..b750bfd5 100644 --- a/src/matricize.jl +++ b/src/matricize.jl @@ -70,10 +70,11 @@ end # ===================================== matricize ======================================== # Copy convention: `bipermutedims`/`permutedims` always copy (Base `permutedims` semantics). # `matricize`/`matricizeperm`/`matricizeopperm` are the maybe-alias tier — the result may be a -# view of the input or a fresh gather, and callers must treat it as read-only. A consumer that -# mutates takes an explicit copy (`MatrixAlgebraKit.copy_input` or `copy`); the factorization -# wrappers donate a provably owned matricization to the mutating matrix-level entries (see -# `maybe_donate` in `factorizations.jl`). +# view of the input or a fresh gather, and callers must treat it as read-only. Write access is +# never inferred from the maybe-alias tier: a consumer that mutates materializes an owned matrix +# (`MatrixAlgebraKit.copy_input` or an always-copy `bipermutedims` gather, see the owned tier in +# `factorizations.jl`) and a consumer that writes into a destination asks the style for a +# memory-sharing matricization (`trymatricizeview`). # This is the primary function that should be overloaded for new matricize styles. # This assumes the permutation was already performed. @@ -135,23 +136,6 @@ function matricizeperm( return matricizeperm(style, a, to_permblocks(a, (perm_codomain, perm_domain))...) end -# `matricizeperm` for the factorization and matrix-function wrappers: same maybe-alias -# contract, but a codomain/domain swap takes the permuted copy instead of the lazy `transpose` -# view — matrix-level backends (LAPACK through MatrixAlgebraKit) require the matricized layout -# itself, not just any matrix-shaped view. -function matricize_input( - style::MatricizeStyle, a, - perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}} - ) - ndims(a) == length(perm_codomain) + length(perm_domain) || - throw(ArgumentError("Invalid bipermutation")) - kind = matricizekind(style, perm_codomain, perm_domain) - kind == ReshapeMatricizeKind && - return matricize(style, a, Val(length(perm_codomain))) - a_perm = bipermutedims(a, perm_codomain, perm_domain) - return matricize(style, a_perm, Val(length(perm_codomain))) -end - # ================================== matricizeopperm ===================================== """ @@ -172,62 +156,39 @@ function matricizeopperm( ) return matricizeopperm(style, op, a, to_permblocks(a, (perm_codomain, perm_domain))...) end -# Classifies how `matricize` realizes the bipermutation `(perm_codomain, perm_domain)` -# against storage, so `matricizeopperm` can skip the redundant permuted copy: -# ReshapeMatricizeKind — the groups are already in storage order, so the permute is a -# no-op and `matricize(style, a, ...)` can be called directly. -# For a dense array that is a `reshape` view; for a graded array -# it still gathers blocks, but skips the extra permute copy. -# TransposeMatricizeKind — the only reordering is a codomain/domain swap, which a dense -# array realizes as a `transpose` of a `reshape` (a view gemm -# reads via BLAS' transpose flag). -# PermuteMatricizeKind — the groups interleave storage, so a permuted copy is required. -# Pure: depends only on the index pattern, not on `a`'s data. Dispatched on `MatricizeStyle`. -# The generic classifier only recognizes the always-safe `ReshapeMatricizeKind` (skipping a -# no-op permute is valid for any style); `TransposeMatricizeKind` is opt-in for styles whose -# `matricize` composes with a lazy `transpose`, currently only `ReshapeMatricize`. -@enum MatricizeKind ReshapeMatricizeKind TransposeMatricizeKind PermuteMatricizeKind - # Whether `perm` is the identity permutation `(1, …, n)`. isidentityperm(perm::Tuple{Vararg{Int}}) = perm == ntuple(identity, length(perm)) -function matricizekind( - ::MatricizeStyle, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}} - ) - # Already in storage order: the permute is a no-op, so `matricize` can run directly. - isidentityperm((perm_codomain..., perm_domain...)) && return ReshapeMatricizeKind - return PermuteMatricizeKind -end - -# Whether `matricizeperm(style, a, perm_codomain, perm_domain)` aliases `a` — returns a view (so a -# `mul!` into it writes through to `a`) rather than freshly allocated storage. Only a style whose -# `matricize` is itself a view (a dense reshape) can alias, and only when the bipermutation needs -# no permuted copy. Defaults to `false`: a style that gathers into new storage, such as a graded -# array, never aliases its input. -matricizepermaliases(::MatricizeStyle, perm_codomain, perm_domain) = false - -# Skip the permuted copy when the classifier says it is unnecessary. `ReshapeMatricizeKind` -# calls `matricize` directly on `a` (a view for dense, a gather without the extra permute -# for graded); `TransposeMatricizeKind` returns a lazy `transpose` of the reshape. Both -# fast paths require `op === identity`, since a plain view cannot carry a fused `op` like -# `conj`. The result may alias `a` and must be treated as read-only, matching the docstring. +# The identity bipermutation is a no-op permute, so `matricize` runs directly on `a` (a view +# for dense, a gather without the extra permute copy for graded); the fast path requires +# `op === identity`, since a plain view cannot carry a fused `op` like `conj`. The result may +# alias `a` and must be treated as read-only, matching the docstring. function matricizeopperm( style::MatricizeStyle, op, a, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}} ) ndims(a) == length(perm_codomain) + length(perm_domain) || throw(ArgumentError("Invalid bipermutation")) - if op === identity - kind = matricizekind(style, perm_codomain, perm_domain) - kind == ReshapeMatricizeKind && - return matricize(style, a, Val(length(perm_codomain))) - kind == TransposeMatricizeKind && - return transpose(matricize(style, a, Val(length(perm_domain)))) - end + op === identity && isidentityperm((perm_codomain..., perm_domain...)) && + return matricize(style, a, Val(length(perm_codomain))) a_perm_op = permutedimsop(op, a, perm_codomain, perm_domain) return matricize(style, a_perm_op, Val(length(perm_codomain))) end +# ================================= trymatricizeview ===================================== +# Return a matricization sharing `a_dest`'s memory (writes to it are writes to `a_dest`), or +# `nothing` when producing the matricization would move data (the try/`nothing` convention of +# `tryflattenlinear`). Styles overload the `Val` (trivial split) form to declare which splits +# share; the bipermutation form delegates to it at the identity. +trymatricizeview(::MatricizeStyle, a_dest, ndims_codomain::Val) = nothing +function trymatricizeview( + style::MatricizeStyle, a_dest, + invperm_codomain::Tuple{Vararg{Int}}, invperm_domain::Tuple{Vararg{Int}} + ) + isidentityperm((invperm_codomain..., invperm_domain...)) || return nothing + return trymatricizeview(style, a_dest, Val(length(invperm_codomain))) +end + # ==================================== unmatricize ======================================= # Split form: `axes_codomain` and `axes_domain` are the destination axes for the codomain and # domain groups, given codomain-facing (un-dualized), the same convention as `similar_map`. A @@ -322,20 +283,21 @@ function matricize(::ReshapeMatricize, a, ndims_codomain::Val) size_codomain, size_domain = bipartition(size(a), ndims_codomain) return reshape(a, (prod(size_codomain), prod(size_domain))) end -# A dense array additionally realizes a codomain/domain swap as a lazy `transpose` of a -# reshape (a view), so it opts into `TransposeMatricizeKind` on top of the generic -# reshape/permute classification. -function matricizekind( - ::ReshapeMatricize, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}} +# A dense reshape matricization is a view of `a_dest` at any order-preserving split, and a +# pure codomain/domain swap is a lazy `transpose` of that view (BLAS consumes it through the +# transpose flag), so both share memory; an interleaving split needs a permuted copy. +function trymatricizeview(style::ReshapeMatricize, a_dest, ndims_codomain::Val) + return matricize(style, a_dest, ndims_codomain) +end +function trymatricizeview( + style::ReshapeMatricize, a_dest, + invperm_codomain::Tuple{Vararg{Int}}, invperm_domain::Tuple{Vararg{Int}} ) - isidentityperm((perm_codomain..., perm_domain...)) && return ReshapeMatricizeKind - isidentityperm((perm_domain..., perm_codomain...)) && return TransposeMatricizeKind - return PermuteMatricizeKind -end -# A dense reshape/transpose is a view of `a`; only a permuted copy detaches. So the matricized -# output aliases `a` for every kind except `PermuteMatricizeKind`. -function matricizepermaliases(style::ReshapeMatricize, perm_codomain, perm_domain) - return matricizekind(style, perm_codomain, perm_domain) != PermuteMatricizeKind + isidentityperm((invperm_codomain..., invperm_domain...)) && + return trymatricizeview(style, a_dest, Val(length(invperm_codomain))) + isidentityperm((invperm_domain..., invperm_codomain...)) && + return transpose(matricize(style, a_dest, Val(length(invperm_domain)))) + return nothing end # The matricized input's rows must be the fused codomain and its columns the fused domain. # `reshape` alone only checks the total element count, so a wrong split with the right total diff --git a/src/matrixfunctions.jl b/src/matrixfunctions.jl index 3d571063..ebbda116 100644 --- a/src/matrixfunctions.jl +++ b/src/matrixfunctions.jl @@ -31,17 +31,17 @@ const MATRIX_FUNCTIONS = [ :acoth, ] -# The wrappers share the factorization machinery: `matricized` cores below consume the -# maybe-alias matricization read-only (the matrix functions allocate their own outputs), so -# the permuted forms skip the eager `bipermutedims` copy (see `factorizations.jl`). +# The matrix functions never mutate their input (they allocate their own outputs), so the +# permuted forms consume the maybe-alias `matricizeperm` matricization read-only, skipping +# the eager `bipermutedims` copy at the identity bipermutation. for f in MATRIX_FUNCTIONS @eval begin - function $f(style::MatricizeStyle, a, ndims_codomain::Val; kwargs...) - a_mat = matricize(style, a, ndims_codomain) - axes_codomain, axes_domain = bipartition_axes(axes(a), ndims_codomain) - return matricized( - $f, style, a_mat, isdetached(a_mat, a), - axes_codomain, axes_domain; kwargs... + function $f(style::MatricizeStyle, a, ndims_codomain::Val{K}; kwargs...) where {K} + return $f( + style, a, + ntuple(identity, ndims_codomain), + ntuple(i -> K + i, Val(ndims(a) - K)); + kwargs... ) end function $f(a, ndims_codomain::Val; kwargs...) @@ -53,15 +53,13 @@ for f in MATRIX_FUNCTIONS perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs... ) - a_mat = matricize_input(style, a, perm_codomain, perm_domain) + a_mat = matricizeperm(style, a, perm_codomain, perm_domain) axes_codomain, axes_domain = bipartition_axes( map(i -> axes(a, i), (perm_codomain..., perm_domain...)), Val(length(perm_codomain)) ) - return matricized( - $f, style, a_mat, isdetached(a_mat, a), - axes_codomain, axes_domain; kwargs... - ) + fa_mat = Base.$f(a_mat; kwargs...) + return unmatricize(style, fa_mat, axes_codomain, axes_domain) end function $f( a, @@ -71,14 +69,6 @@ for f in MATRIX_FUNCTIONS return $f(MatricizeStyle(a), a, perm_codomain, perm_domain; kwargs...) end - function matricized( - ::typeof($f), style::MatricizeStyle, a_mat, owned::Bool, - axes_codomain, axes_domain; kwargs... - ) - fa_mat = Base.$f(a_mat; kwargs...) - return unmatricize(style, fa_mat, axes_codomain, axes_domain) - end - function $f( style::MatricizeStyle, a, labels_a, labels_codomain, labels_domain; kwargs... diff --git a/test/test_factorizations.jl b/test/test_factorizations.jl index 45357206..f8934459 100644 --- a/test/test_factorizations.jl +++ b/test/test_factorizations.jl @@ -419,44 +419,102 @@ end @test TensorAlgebra.tr(A, Val(2)) ≈ LinearAlgebra.tr(m) end -# Permuted entry points: matricization and buffer donation -# -------------------------------------------------------- -# The permuted forms matricize directly (`matricize_input`, no eager `bipermutedims` copy) -# and donate an owned matricization to the mutating MatrixAlgebraKit entries, so pin them -# to the reference permute-then-`Val` path and check the caller's array is never mutated, -# through the identity, permuted-copy, and codomain/domain-swap matricizations alike. -@testset "Permuted forms match permute-then-matricize ($T)" for T in elts +# Permuted entry points +# --------------------- +# The permuted entry points must not mutate the caller's array, through the identity, +# permuted-copy, and codomain/domain-swap matricizations alike. +@testset "Permuted forms: input preserved, factors reconstruct ($T)" for T in elts A = randn(T, 2, 3, 4) Acopy = copy(A) for (perm_codomain, perm_domain) in (((1, 2), (3,)), ((3, 1), (2,)), ((3,), (1, 2)), ((2,), (3, 1))) + k = length(perm_codomain) A_perm = TensorAlgebra.bipermutedims(A, perm_codomain, perm_domain) - for f in ( - qr_compact, lq_compact, left_orth, right_orth, - svd_compact, svd_trunc, svd_vals, left_null, right_null, - ) - F = f(A, perm_codomain, perm_domain) - F_ref = f(A_perm, Val(length(perm_codomain))) - Fs = F isa Tuple ? F : (F,) - F_refs = F_ref isa Tuple ? F_ref : (F_ref,) - @test all(map(==, Fs, F_refs)) - @test A == Acopy + A_mat = TensorAlgebra.matricize(A_perm, Val(k)) + for f in (qr_compact, lq_compact, left_orth, right_orth) + X, Y = f(A, perm_codomain, perm_domain) + @test TensorAlgebra.matricize(X, Val(k)) * + TensorAlgebra.matricize(Y, Val(1)) ≈ A_mat + end + for f in (svd_compact, svd_trunc) + U, S, Vᴴ = f(A, perm_codomain, perm_domain) + U_mat = TensorAlgebra.matricize(U, Val(k)) + @test U_mat * S * TensorAlgebra.matricize(Vᴴ, Val(1)) ≈ A_mat + @test U_mat' * U_mat ≈ I end + @test svd_vals(A, perm_codomain, perm_domain) ≈ LinearAlgebra.svdvals(A_mat) + N = TensorAlgebra.matricize(left_null(A, perm_codomain, perm_domain), Val(k)) + @test norm(N' * A_mat) ≈ 0 atol = 1.0e-13 + @test N' * N ≈ I + Nᴴ = TensorAlgebra.matricize(right_null(A, perm_codomain, perm_domain), Val(1)) + @test norm(A_mat * Nᴴ') ≈ 0 atol = 1.0e-13 + @test Nᴴ * Nᴴ' ≈ I + @test A == Acopy end B = randn(T, 2, 3, 2, 3) Bcopy = copy(B) for (perm_codomain, perm_domain) in (((1, 2), (3, 4)), ((3, 4), (1, 2)), ((2, 3), (4, 1))) B_perm = TensorAlgebra.bipermutedims(B, perm_codomain, perm_domain) - for f in (eig_full, eig_vals) - F = f(B, perm_codomain, perm_domain) - F_ref = f(B_perm, Val(2)) - Fs = F isa Tuple ? F : (F,) - F_refs = F_ref isa Tuple ? F_ref : (F_ref,) - @test all(map(==, Fs, F_refs)) - @test B == Bcopy - end - @test TensorAlgebra.tr(B, perm_codomain, perm_domain) ≈ - TensorAlgebra.tr(B_perm, Val(2)) + B_mat = Matrix(TensorAlgebra.matricize(B_perm, Val(2))) + D, V = eig_full(B, perm_codomain, perm_domain) + V_mat = TensorAlgebra.matricize(V, Val(2)) + @test B_mat * V_mat ≈ V_mat * D + sortvals(v) = sort(v; by = x -> (real(x), imag(x))) + @test sortvals(eig_vals(B, perm_codomain, perm_domain)) ≈ + sortvals(LinearAlgebra.eigvals(B_mat)) + @test TensorAlgebra.tr(B, perm_codomain, perm_domain) ≈ LinearAlgebra.tr(B_mat) + @test B == Bcopy + end +end + +# A wrapper whose `matricize` reshapes the parent's buffer but declares nothing about it (no +# `Base.dataids` overload), so any ownership inference from aliasing checks misclassifies it. +module FactorizationMatricizeTestUtils + using TensorAlgebra: TensorAlgebra as TA + struct AliasingArray{T, N, P <: AbstractArray{T, N}} <: AbstractArray{T, N} + parent::P + end + Base.size(a::AliasingArray) = size(a.parent) + function Base.getindex(a::AliasingArray{<:Any, N}, I::Vararg{Int, N}) where {N} + return a.parent[I...] end + struct AliasingMatricize <: TA.MatricizeStyle end + TA.MatricizeStyle(::Type{<:AliasingArray}) = AliasingMatricize() + function TA.matricize(::AliasingMatricize, a::AliasingArray, ndims_codomain::Val) + return TA.matricize(TA.ReshapeMatricize(), a.parent, ndims_codomain) + end + function TA.matricize(::AliasingMatricize, a::AbstractArray, ndims_codomain::Val) + return TA.matricize(TA.ReshapeMatricize(), a, ndims_codomain) + end + function TA.unmatricize(::AliasingMatricize, m, axes_codomain, axes_domain) + return AliasingArray( + TA.unmatricize(TA.ReshapeMatricize(), m, axes_codomain, axes_domain) + ) + end +end +using .FactorizationMatricizeTestUtils: AliasingArray + +@testset "Aliasing matricize wrapper: input preserved ($f)" for f in + (qr_compact, svd_compact) + parent = randn(2, 3, 4) + A = AliasingArray(parent) + parent_copy = copy(parent) + f(A, Val(1)) + @test parent == parent_copy + f(A, (1,), (2, 3)) + @test parent == parent_copy +end + +@testset "Integer eltype through the wrappers" begin + parent = rand(-9:9, 2, 3, 4) + A = AliasingArray(parent) + parent_copy = copy(parent) + Q, R = qr_compact(A, Val(1)) + @test parent == parent_copy + Q_mat = TensorAlgebra.matricize(Q, Val(1)) + @test eltype(Q_mat) === Float64 + @test Q_mat * TensorAlgebra.matricize(R, Val(1)) ≈ reshape(parent, 2, 12) + @test svd_vals(A, (1,), (2, 3)) ≈ LinearAlgebra.svdvals(reshape(float.(parent), 2, 12)) + @test parent == parent_copy end diff --git a/test/test_matricize.jl b/test/test_matricize.jl index bb9139a1..4402bd93 100644 --- a/test/test_matricize.jl +++ b/test/test_matricize.jl @@ -1,8 +1,7 @@ -using LinearAlgebra: Transpose using StableRNGs: StableRNG -using TensorAlgebra: TensorAlgebra, PermuteMatricizeKind, ReshapeMatricize, - ReshapeMatricizeKind, TransposeMatricizeKind, matricizeopperm, matricizeperm -using Test: @test, @testset +using TensorAlgebra: + TensorAlgebra, ReshapeMatricize, matricizeopperm, matricizeperm, trymatricizeview +using Test: @test, @test_throws, @testset # A non-`ReshapeMatricize` style, to check the always-safe generic fallback. struct DummyMatricize <: TensorAlgebra.MatricizeStyle end @@ -15,52 +14,25 @@ function matricize_ref(a, perm_codomain, perm_domain) return reshape(a_perm, (nrow, ncol)) end -@testset "matricizekind classifier" begin - style = ReshapeMatricize() - # Already in storage order → plain reshape view. - @test TensorAlgebra.matricizekind(style, (1,), (2, 3)) == ReshapeMatricizeKind - @test TensorAlgebra.matricizekind(style, (1, 2), (3,)) == ReshapeMatricizeKind - @test TensorAlgebra.matricizekind(style, (1, 2, 3), ()) == ReshapeMatricizeKind - @test TensorAlgebra.matricizekind(style, (), (1, 2, 3)) == ReshapeMatricizeKind - # Pure codomain/domain swap → transpose of a reshape view. - @test TensorAlgebra.matricizekind(style, (2, 3), (1,)) == TransposeMatricizeKind - @test TensorAlgebra.matricizekind(style, (3,), (1, 2)) == TransposeMatricizeKind - # Interleaved → permuted copy. - @test TensorAlgebra.matricizekind(style, (3, 1), (2,)) == PermuteMatricizeKind - @test TensorAlgebra.matricizekind(style, (2,), (1, 3)) == PermuteMatricizeKind - @test TensorAlgebra.matricizekind(style, (1, 3), (2,)) == PermuteMatricizeKind - # Generic matricize styles recognize the always-safe reshape (no-op permute) but never - # claim a transpose (which only styles with a lazy `transpose` can realize). - @test TensorAlgebra.matricizekind(DummyMatricize(), (1,), (2, 3)) == - ReshapeMatricizeKind - @test TensorAlgebra.matricizekind(DummyMatricize(), (1, 2, 3), ()) == - ReshapeMatricizeKind - @test TensorAlgebra.matricizekind(DummyMatricize(), (2, 3), (1,)) == - PermuteMatricizeKind - @test TensorAlgebra.matricizekind(DummyMatricize(), (3, 1), (2,)) == - PermuteMatricizeKind -end - -@testset "maybe-view matricizeopperm (eltype=$elt)" for elt in (Float64, ComplexF64) +@testset "maybe-view matricizeperm (eltype=$elt)" for elt in (Float64, ComplexF64) a = randn(StableRNG(123), elt, 2, 3, 4) - # Reshape branch: correct values and a view aliasing `a`. + # Identity bipermutation: correct values and a view aliasing `a`. m = matricizeperm(a, (1,), (2, 3)) @test m ≈ matricize_ref(a, (1,), (2, 3)) @test Base.mightalias(m, a) - # Transpose branch: correct values and a transpose view aliasing `a`. - m = matricizeperm(a, (2, 3), (1,)) - @test m ≈ matricize_ref(a, (2, 3), (1,)) - @test m isa Transpose - @test Base.mightalias(m, a) - - # Permute branch: correct values, but a fresh copy (no aliasing). - m = matricizeperm(a, (3, 1), (2,)) - @test m ≈ matricize_ref(a, (3, 1), (2,)) - @test !Base.mightalias(m, a) + # Every other bipermutation is a fresh permuted copy in matricized layout (no lazy + # wrappers), including the codomain/domain swap. + for (pc, pd) in (((2, 3), (1,)), ((3, 1), (2,))) + m = matricizeperm(a, pc, pd) + @test m ≈ matricize_ref(a, pc, pd) + @test m isa Matrix + @test !Base.mightalias(m, a) + end + @test_throws ArgumentError matricizeperm(a, (1,), (2,)) - # `conj` cannot ride a view, so it copies even on the reshape/transpose patterns. + # `conj` cannot ride a view, so it copies even on the identity bipermutation. m = matricizeopperm(conj, a, (1,), (2, 3)) @test m ≈ conj.(matricize_ref(a, (1,), (2, 3))) @test !Base.mightalias(m, a) @@ -69,25 +41,48 @@ end @test !Base.mightalias(m, a) end -@testset "view branches track source mutations, copy branch does not" begin +@testset "trymatricizeview" begin + a = randn(StableRNG(321), 2, 3, 4) + style = ReshapeMatricize() + + # Order-preserving splits share memory: a reshape view for the identity bipermutation, + # a lazy transpose of it for the codomain/domain swap. + m = trymatricizeview(style, a, (1,), (2, 3)) + @test m == matricize_ref(a, (1,), (2, 3)) + @test Base.mightalias(m, a) + @test trymatricizeview(style, a, Val(1)) == m + m = trymatricizeview(style, a, (2, 3), (1,)) + @test m == matricize_ref(a, (2, 3), (1,)) + @test Base.mightalias(m, a) + + # An interleaving split would move data, so no memory-sharing matricization exists. + @test isnothing(trymatricizeview(style, a, (3, 1), (2,))) + + # A generic style declares nothing. + @test isnothing(trymatricizeview(DummyMatricize(), a, Val(1))) + @test isnothing(trymatricizeview(DummyMatricize(), a, (1,), (2, 3))) + + # Writes to the shared matricization are writes to `a`. + m = trymatricizeview(style, a, (1,), (2, 3)) + m[1, 1] = 42 + @test a[1, 1, 1] == 42 +end + +@testset "view branch tracks source mutations, copy branch does not" begin rng = StableRNG(7) - # Reshape view tracks an in-place update of `a`. + # Identity-bipermutation view tracks an in-place update of `a`. a = randn(rng, 2, 3, 4) m = matricizeperm(a, (1,), (2, 3)) a .= randn(rng, 2, 3, 4) @test m ≈ matricize_ref(a, (1,), (2, 3)) - # Transpose view tracks an in-place update of `a`. - a = randn(rng, 2, 3, 4) - m = matricizeperm(a, (2, 3), (1,)) - a .= randn(rng, 2, 3, 4) - @test m ≈ matricize_ref(a, (2, 3), (1,)) - - # Permute copy is independent of later updates to `a`. - a = randn(rng, 2, 3, 4) - m = matricizeperm(a, (3, 1), (2,)) - snapshot = copy(m) - a .= a .+ 1 - @test m == snapshot + # Permuted copies are independent of later updates to `a`. + for (pc, pd) in (((2, 3), (1,)), ((3, 1), (2,))) + a = randn(rng, 2, 3, 4) + m = matricizeperm(a, pc, pd) + snapshot = copy(m) + a .= a .+ 1 + @test m == snapshot + end end diff --git a/test/test_matrixalgebra.jl b/test/test_matrixalgebra.jl index af75ebd1..291db7a8 100644 --- a/test/test_matrixalgebra.jl +++ b/test/test_matrixalgebra.jl @@ -2,7 +2,7 @@ using LinearAlgebra: Diagonal, I, diag, norm using MatrixAlgebraKit: qr_compact, svd_trunc, truncrank using StableRNGs: StableRNG using TensorAlgebra.MatrixAlgebra: MatrixAlgebra, truncdegen -using Test: @test, @testset +using Test: @test, @test_throws, @testset elts = (Float32, Float64, ComplexF32, ComplexF64) From cc8c729061840698472573e62367ad63707668b5 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Thu, 10 Sep 2026 14:33:53 -0400 Subject: [PATCH 3/5] Route matricize through an ismatricizeview trait Co-Authored-By: Claude Fable 5 --- ext/TensorAlgebraTensorKitExt.jl | 22 +++++- src/contract/contract_matricize.jl | 6 +- src/diagonal.jl | 5 +- src/factorizations.jl | 20 +++--- src/matricize.jl | 105 ++++++++++++++++++----------- test/test_matricize.jl | 86 ++++++++++++++++++----- test/test_matricizestyle.jl | 26 +++++++ 7 files changed, 194 insertions(+), 76 deletions(-) diff --git a/ext/TensorAlgebraTensorKitExt.jl b/ext/TensorAlgebraTensorKitExt.jl index 40075a08..aa770db2 100644 --- a/ext/TensorAlgebraTensorKitExt.jl +++ b/ext/TensorAlgebraTensorKitExt.jl @@ -243,11 +243,29 @@ end struct TensorKitMatricize <: TensorAlgebra.MatricizeStyle end TensorAlgebra.MatricizeStyle(::Type{<:AbstractTensorMap}) = TensorKitMatricize() -function TensorAlgebra.matricize( +# `permute` at the tensor's own codomain/domain split is trivial and returns `t` itself, so +# the matching split is the one memory-sharing matricization (TensorKit's own +# `has_shared_permute` notion); any other split regroups into a fresh `TensorMap`. +function TensorAlgebra.ismatricizeview( + ::TensorKitMatricize, ::AbstractTensorMap{<:Any, <:Any, K}, ::Val{K} + ) where {K} + return true +end +TensorAlgebra.ismatricizeview(::TensorKitMatricize, ::AbstractTensorMap, ::Val) = false +function TensorAlgebra.matricizeview( + ::TensorKitMatricize, t::AbstractTensorMap{<:Any, <:Any, K}, ::Val{K} + ) where {K} + return t +end +function TensorAlgebra.matricizecopy( ::TensorKitMatricize, t::AbstractTensorMap, ndims_codomain::Val{K} ) where {K} N = numind(t) - return permute(t, (ntuple(identity, Val(K)), ntuple(i -> K + i, Val(N - K)))) + return permute( + t, + (ntuple(identity, Val(K)), ntuple(i -> K + i, Val(N - K))); + copy = true + ) end # The identity fill on the regrouped map is TensorKit's own `one!` (MatrixAlgebraKit's diff --git a/src/contract/contract_matricize.jl b/src/contract/contract_matricize.jl index a1ff5a5f..9e04f5f0 100644 --- a/src/contract/contract_matricize.jl +++ b/src/contract/contract_matricize.jl @@ -23,9 +23,9 @@ function contractopadd!( algorithm.right_matricize_style, op2, a2, biperm2_codomain, biperm2_domain ) output_style = algorithm.output_matricize_style - a_dest_mat = trymatricizeview(output_style, a_dest, invperm_codomain, invperm_domain) - if !isnothing(a_dest_mat) + if ismatricizeview(output_style, a_dest, invperm_codomain, invperm_domain) # The matricization shares `a_dest`'s memory, so the matmul is the whole operation. + a_dest_mat = matricizeview(output_style, a_dest, Val(length(invperm_codomain))) mul!(a_dest_mat, a1_mat, a2_mat, α, β) elseif iszero(β) # `β` is a strong zero, so `a_dest`'s current data is irrelevant: let the matmul @@ -38,7 +38,7 @@ function contractopadd!( else # `a_dest`'s data contributes through `β`, so gather it, multiply into the gathered # copy, and scatter back. - a_dest_mat = matricizeperm(output_style, a_dest, invperm_codomain, invperm_domain) + a_dest_mat = matricizecopy(output_style, a_dest, invperm_codomain, invperm_domain) mul!(a_dest_mat, a1_mat, a2_mat, α, β) unmatricizeperm!(output_style, a_dest, a_dest_mat, invperm_codomain, invperm_domain) end diff --git a/src/diagonal.jl b/src/diagonal.jl index bbb62cd3..7a465893 100644 --- a/src/diagonal.jl +++ b/src/diagonal.jl @@ -42,8 +42,9 @@ function allocate_output( end # A `Diagonal` is already a matrix; the `(1 codomain, 1 domain)` matricization is the identity -# reshape, so return it directly (maybe-alias, matching `matricize`'s general contract). -matricize(::ReshapeMatricize, a::Diagonal, ::Val{1}) = a +# reshape, so the memory-sharing matricization is `a` itself (keeping it a `Diagonal` for the +# `Diagonal`-specialized consumers downstream). +matricizeview(::ReshapeMatricize, a::Diagonal, ::Val{1}) = a # A `{1,1}` unmatricize (one codomain axis, one domain axis) is the endomorphism identity: the # result stays `Diagonal`, so return `m` directly. The generic `check_input(unmatricize, ...)` # validates the axis lengths against `m`'s size. diff --git a/src/factorizations.jl b/src/factorizations.jl index 7bdf53f2..7a98828e 100644 --- a/src/factorizations.jl +++ b/src/factorizations.jl @@ -17,7 +17,7 @@ using MatrixAlgebraKit: MatrixAlgebraKit # Owned tier: the matrix-level entries mutate their input, so the perm form materializes an # owned matricization following MatrixAlgebraKit's `f(A) = f!(copy_input(f, A))` convention — # a memory-sharing matricization is materialized through `MatrixAlgebraKit.copy_input`, while -# the gathered bipermutation is fresh by construction and is donated directly (with +# the `matricizecopy` gather is owned by contract and is donated directly (with # `copy_input` still applied when the eltype must change) — and the cores call the mutating # entry unconditionally. for f in ( @@ -35,12 +35,11 @@ for f in ( ) ndims(A) == length(perm_codomain) + length(perm_domain) || throw(ArgumentError("Invalid bipermutation")) - A_shared = trymatricizeview(style, A, perm_codomain, perm_domain) - A_mat = if !isnothing(A_shared) + A_mat = if ismatricizeview(style, A, perm_codomain, perm_domain) + A_shared = matricizeview(style, A, Val(length(perm_codomain))) MatrixAlgebraKit.copy_input(MatrixAlgebraKit.$f, A_shared) else - A_perm = bipermutedims(A, perm_codomain, perm_domain) - A_gather = matricize(style, A_perm, Val(length(perm_codomain))) + A_gather = matricizecopy(style, A, perm_codomain, perm_domain) if eltype(A_gather) === float(eltype(A_gather)) A_gather else @@ -882,18 +881,17 @@ function one!!(A, ndims_codomain::Val; kwargs...) return one!!(MatricizeStyle(A), A, ndims_codomain; kwargs...) end -# In-place identity fill: writes the identity into `A` and returns it. Fills a memory-sharing +# In-place identity fill: writes the identity into `A` and returns it. Fills the memory-sharing # matricization directly when the style declares one at this split, and otherwise fills a # gathered matrix and scatters it back with `unmatricize!`. function one!(style::MatricizeStyle, A, ndims_codomain::Val; kwargs...) - A_mat = trymatricizeview(style, A, ndims_codomain) - if !isnothing(A_mat) - MatrixAlgebraKit.one!(A_mat) + if ismatricizeview(style, A, ndims_codomain) + MatrixAlgebraKit.one!(matricizeview(style, A, ndims_codomain)) return A end - A_mat = matricize(style, A, ndims_codomain) + A_mat = matricizecopy(style, A, ndims_codomain) MatrixAlgebraKit.one!(A_mat) - return unmatricize!(A, A_mat, ndims_codomain) + return unmatricize!(style, A, A_mat, ndims_codomain) end function one!(A, ndims_codomain::Val; kwargs...) return one!(MatricizeStyle(A), A, ndims_codomain; kwargs...) diff --git a/src/matricize.jl b/src/matricize.jl index b750bfd5..972cb072 100644 --- a/src/matricize.jl +++ b/src/matricize.jl @@ -36,6 +36,11 @@ Non-mutating version of `bipermutedimsopadd!`: returns `op.(permutedims(src, (perm_codomain..., perm_domain...)))`. """ function permutedimsop(op, src, perm_codomain, perm_domain) + # Validate against `src` here: `bipermutedimsopadd!`'s `check_input` compares against `dest`, + # which `allocate_output` builds from the same perms, so it cannot catch a non-covering perm. + perm = (perm_codomain..., perm_domain...) + (ndims(src) == length(perm) && isperm(perm)) || + throw(ArgumentError("Invalid bipermutation")) dest = allocate_output(permutedimsop, op, src, perm_codomain, perm_domain) return bipermutedimsopadd!(dest, op, src, perm_codomain, perm_domain, true, false) end @@ -68,25 +73,49 @@ function bipermutedims!( end # ===================================== matricize ======================================== -# Copy convention: `bipermutedims`/`permutedims` always copy (Base `permutedims` semantics). -# `matricize`/`matricizeperm`/`matricizeopperm` are the maybe-alias tier — the result may be a -# view of the input or a fresh gather, and callers must treat it as read-only. Write access is -# never inferred from the maybe-alias tier: a consumer that mutates materializes an owned matrix -# (`MatrixAlgebraKit.copy_input` or an always-copy `bipermutedims` gather, see the owned tier in -# `factorizations.jl`) and a consumer that writes into a destination asks the style for a -# memory-sharing matricization (`trymatricizeview`). +# Copy convention: `bipermutedims`/`permutedims` always copy (Base `permutedims` semantics). At +# the trivial (`Val`) split the sharing story is exact: `matricizeview` shares `a`'s memory, +# `matricizecopy` returns fresh storage the caller owns, and `matricize` aliases `a` iff +# `ismatricizeview` — so a consumer that writes into a destination checks the trait and writes +# through `matricizeview`, and a consumer that mutates an input either owns a `matricizecopy` +# result by contract or materializes an owned matrix with `MatrixAlgebraKit.copy_input` (see the +# owned tier in `factorizations.jl`). `matricizeperm`/`matricizeopperm` keep maybe-alias +# semantics (the result may view or copy; treat it as read-only) until the planned op/perm-form +# trait lands, and `matricizeview` deliberately has no perm form pending that op/perm-layer +# design. -# This is the primary function that should be overloaded for new matricize styles. -# This assumes the permutation was already performed. -function matricize( - style::MatricizeStyle, a, ndims_codomain::Val - ) - return throw(MethodError(matricize, (style, a, ndims_codomain))) +# `matricize` at the trivial split routes on the style's sharing declaration. Styles implement +# the three leaves (`ismatricizeview`, `matricizeview`, `matricizecopy`) rather than overloading +# `matricize` itself. This assumes the permutation was already performed. +function matricize(style::MatricizeStyle, a, ndims_codomain::Val) + ismatricizeview(style, a, ndims_codomain) && + return matricizeview(style, a, ndims_codomain) + return matricizecopy(style, a, ndims_codomain) end function matricize(a, ndims_codomain::Val) return matricize(MatricizeStyle(a), a, ndims_codomain) end +# Partial: defined only where `ismatricizeview` is `true`, and always returns a matricization +# sharing `a`'s memory (the `StridedView` partial-constructor pattern). +function matricizeview(style::MatricizeStyle, a, ndims_codomain::Val) + return throw(MethodError(matricizeview, (style, a, ndims_codomain))) +end +# Total: always returns a matricization in fresh storage the caller owns. +function matricizecopy(style::MatricizeStyle, a, ndims_codomain::Val) + return throw(MethodError(matricizecopy, (style, a, ndims_codomain))) +end + +# `bipermutedims` always copies and `matricize` might return a view, so the result is +# guaranteed to be a copy. +function matricizecopy( + style::MatricizeStyle, a, + perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}} + ) + a_perm = bipermutedims(a, perm_codomain, perm_domain) + return matricize(style, a_perm, Val(length(perm_codomain))) +end + function matricizeperm( a, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}} @@ -175,18 +204,21 @@ function matricizeopperm( return matricize(style, a_perm_op, Val(length(perm_codomain))) end -# ================================= trymatricizeview ===================================== -# Return a matricization sharing `a_dest`'s memory (writes to it are writes to `a_dest`), or -# `nothing` when producing the matricization would move data (the try/`nothing` convention of -# `tryflattenlinear`). Styles overload the `Val` (trivial split) form to declare which splits -# share; the bipermutation form delegates to it at the identity. -trymatricizeview(::MatricizeStyle, a_dest, ndims_codomain::Val) = nothing -function trymatricizeview( - style::MatricizeStyle, a_dest, +# ================================== ismatricizeview ===================================== +# `true` iff `matricize(style, a, ndims_codomain)` shares `a`'s memory (writes to it are writes +# to `a`) — the `isstrided`/`StridedView` pattern (also TensorKit's `has_shared_permute` and +# TensorOperations' `isblasdestination`). Styles overload the `Val` (trivial split) form to +# declare which splits share; the bipermutation form delegates to it at the identity and is +# `false` (fail-safe) everywhere else. A general `ismatricizeview(style, op, a, perm_codomain, +# perm_domain)` form (op and bipermutation view-sets) is planned; these are its +# `op === identity` special cases. +ismatricizeview(::MatricizeStyle, a, ndims_codomain::Val) = false +function ismatricizeview( + style::MatricizeStyle, a, invperm_codomain::Tuple{Vararg{Int}}, invperm_domain::Tuple{Vararg{Int}} ) - isidentityperm((invperm_codomain..., invperm_domain...)) || return nothing - return trymatricizeview(style, a_dest, Val(length(invperm_codomain))) + isidentityperm((invperm_codomain..., invperm_domain...)) || return false + return ismatricizeview(style, a, Val(length(invperm_codomain))) end # ==================================== unmatricize ======================================= @@ -263,41 +295,34 @@ end # scatter the fused matrix `m` back into `a_dest`'s existing storage across the codomain/domain # split at `ndims_codomain`. The split applies no permutation, so this is `unmatricizeperm!` at the # trivial bipermutation, reusing its in-place block scatter (no intermediate `unmatricize` copy). -function unmatricize!(a_dest, m, ndims_codomain::Val) +function unmatricize!(style::MatricizeStyle, a_dest, m, ndims_codomain::Val) K = unval(ndims_codomain) N = ndims(a_dest) return unmatricizeperm!( + style, a_dest, m, ntuple(identity, Val(K)), ntuple(i -> K + i, Val(N - K)) ) end +function unmatricize!(a_dest, m, ndims_codomain::Val) + return unmatricize!(MatricizeStyle(a_dest), a_dest, m, ndims_codomain) +end # Defaults to ReshapeMatricize, a simple reshape struct ReshapeMatricize <: MatricizeStyle end MatricizeStyle(::Type{<:AbstractArray}) = ReshapeMatricize() -function matricize(::ReshapeMatricize, a, ndims_codomain::Val) +# A dense reshape matricization is a lazy wrapper at any split, so it always shares memory. +ismatricizeview(::ReshapeMatricize, a, ndims_codomain::Val) = true +function matricizeview(::ReshapeMatricize, a, ndims_codomain::Val) unval(ndims_codomain) ≤ ndims(a) || throw(ArgumentError("Codomain length exceeds number of dimensions.")) size_codomain, size_domain = bipartition(size(a), ndims_codomain) return reshape(a, (prod(size_codomain), prod(size_domain))) end -# A dense reshape matricization is a view of `a_dest` at any order-preserving split, and a -# pure codomain/domain swap is a lazy `transpose` of that view (BLAS consumes it through the -# transpose flag), so both share memory; an interleaving split needs a permuted copy. -function trymatricizeview(style::ReshapeMatricize, a_dest, ndims_codomain::Val) - return matricize(style, a_dest, ndims_codomain) -end -function trymatricizeview( - style::ReshapeMatricize, a_dest, - invperm_codomain::Tuple{Vararg{Int}}, invperm_domain::Tuple{Vararg{Int}} - ) - isidentityperm((invperm_codomain..., invperm_domain...)) && - return trymatricizeview(style, a_dest, Val(length(invperm_codomain))) - isidentityperm((invperm_domain..., invperm_codomain...)) && - return transpose(matricize(style, a_dest, Val(length(invperm_domain)))) - return nothing +function matricizecopy(style::ReshapeMatricize, a, ndims_codomain::Val) + return copy(matricizeview(style, a, ndims_codomain)) end # The matricized input's rows must be the fused codomain and its columns the fused domain. # `reshape` alone only checks the total element count, so a wrong split with the right total diff --git a/test/test_matricize.jl b/test/test_matricize.jl index 4402bd93..d16eba71 100644 --- a/test/test_matricize.jl +++ b/test/test_matricize.jl @@ -1,6 +1,6 @@ using StableRNGs: StableRNG -using TensorAlgebra: - TensorAlgebra, ReshapeMatricize, matricizeopperm, matricizeperm, trymatricizeview +using TensorAlgebra: TensorAlgebra, ReshapeMatricize, ismatricizeview, matricize, + matricizecopy, matricizeopperm, matricizeperm, matricizeview using Test: @test, @test_throws, @testset # A non-`ReshapeMatricize` style, to check the always-safe generic fallback. @@ -41,33 +41,83 @@ end @test !Base.mightalias(m, a) end -@testset "trymatricizeview" begin +@testset "ismatricizeview" begin a = randn(StableRNG(321), 2, 3, 4) style = ReshapeMatricize() - # Order-preserving splits share memory: a reshape view for the identity bipermutation, - # a lazy transpose of it for the codomain/domain swap. - m = trymatricizeview(style, a, (1,), (2, 3)) - @test m == matricize_ref(a, (1,), (2, 3)) - @test Base.mightalias(m, a) - @test trymatricizeview(style, a, Val(1)) == m - m = trymatricizeview(style, a, (2, 3), (1,)) - @test m == matricize_ref(a, (2, 3), (1,)) - @test Base.mightalias(m, a) + # A dense reshape matricization shares memory at every trivial split. + @test ismatricizeview(style, a, Val(1)) + @test ismatricizeview(style, a, (1,), (2, 3)) - # An interleaving split would move data, so no memory-sharing matricization exists. - @test isnothing(trymatricizeview(style, a, (3, 1), (2,))) + # The bipermutation form declares sharing only at the identity: a swap or interleaving + # bipermutation routes through the consumers' gather branches. + @test !ismatricizeview(style, a, (2, 3), (1,)) + @test !ismatricizeview(style, a, (3, 1), (2,)) - # A generic style declares nothing. - @test isnothing(trymatricizeview(DummyMatricize(), a, Val(1))) - @test isnothing(trymatricizeview(DummyMatricize(), a, (1,), (2, 3))) + # A generic style declares nothing (fail-safe default). + @test !ismatricizeview(DummyMatricize(), a, Val(1)) + @test !ismatricizeview(DummyMatricize(), a, (1,), (2, 3)) # Writes to the shared matricization are writes to `a`. - m = trymatricizeview(style, a, (1,), (2, 3)) + m = matricizeview(style, a, Val(1)) + @test m == matricize_ref(a, (1,), (2, 3)) m[1, 1] = 42 @test a[1, 1, 1] == 42 end +@testset "ismatricizeview coherence" begin + rng = StableRNG(11) + a = randn(rng, 2, 3, 4) + style = ReshapeMatricize() + + # A declared share means `matricizeview` (and so `matricize`) aliases `a`, while + # `matricizecopy` never does. + for K in 0:3 + if ismatricizeview(style, a, Val(K)) + m = matricizeview(style, a, Val(K)) + @test Base.mightalias(m, a) + @test matricize(style, a, Val(K)) == m + end + @test !Base.mightalias(matricizecopy(style, a, Val(K)), a) + + # The perm form of the copy leaf is owned too, and at the trivial bipermutation it + # matches the `Val` form. + pc = ntuple(identity, K) + pd = ntuple(i -> K + i, 3 - K) + m_perm = matricizecopy(style, a, pc, pd) + @test !Base.mightalias(m_perm, a) + @test m_perm == matricizecopy(style, a, Val(K)) + end + for (pc, pd) in (((2, 3), (1,)), ((3, 1), (2,))) + m = matricizecopy(style, a, pc, pd) + @test m ≈ matricize_ref(a, pc, pd) + @test !Base.mightalias(m, a) + end + @test_throws ArgumentError matricizecopy(style, a, (1,), (2,)) + + # Both destination branches of a consumer (`contractadd!`) behave: the shared-view route + # for the identity destination bipermutation and the gather/scatter route otherwise. + a1 = randn(rng, 2, 3, 5) + a2 = randn(rng, 5, 3, 2) + ref = TensorAlgebra.contract((:i, :j, :k, :l), a1, (:i, :j, :m), a2, (:m, :k, :l)) + for labels in ((:i, :j, :k, :l), (:k, :l, :i, :j), (:k, :i, :l, :j)) + perm = map(l -> findfirst(==(l), (:i, :j, :k, :l)), labels) + dest = randn(rng, map(d -> size(ref, d), perm)...) + expected = permutedims(ref, perm) .+ 2.0 .* dest + TensorAlgebra.contractadd!( + dest, + labels, + a1, + (:i, :j, :m), + a2, + (:m, :k, :l), + 1.0, + 2.0 + ) + @test dest ≈ expected + end +end + @testset "view branch tracks source mutations, copy branch does not" begin rng = StableRNG(7) diff --git a/test/test_matricizestyle.jl b/test/test_matricizestyle.jl index 9ff7dd37..c43627f1 100644 --- a/test/test_matricizestyle.jl +++ b/test/test_matricizestyle.jl @@ -1,3 +1,4 @@ +using LinearAlgebra: I using TensorAlgebra: TensorAlgebra as TA, Matricize, MatricizeStyle, ReshapeMatricize using Test: @test, @testset @@ -8,6 +9,22 @@ module MatricizeStyleTestUtils end struct MyArrayMatricize <: TA.MatricizeStyle end TA.MatricizeStyle(::Type{<:MyArray}) = MyArrayMatricize() + # Minimal fold/unfold leaves so a round-trip (`one!`) can run through the custom style: + # both dispatch on `MyArrayMatricize`, so an unfold whose style was re-derived from the + # plain fused matrix instead of threaded through would miss them and error. + TA.ismatricizeview(::MyArrayMatricize, a, ::Val) = false + function TA.matricizecopy(::MyArrayMatricize, a::MyArray, ndims_codomain::Val) + return TA.matricizecopy(TA.ReshapeMatricize(), a.parent, ndims_codomain) + end + function TA.unmatricizeperm!( + ::MyArrayMatricize, a_dest::MyArray, m, + invperm_codomain::Tuple{Vararg{Int}}, invperm_domain::Tuple{Vararg{Int}} + ) + TA.unmatricizeperm!( + TA.ReshapeMatricize(), a_dest.parent, m, invperm_codomain, invperm_domain + ) + return a_dest + end end using .MatricizeStyleTestUtils: MyArray, MyArrayMatricize @@ -30,3 +47,12 @@ using .MatricizeStyleTestUtils: MyArray, MyArrayMatricize @test TA.default_contract_algorithm(typeof(a2), typeof(a2)) ≡ Matricize(MyArrayMatricize()) end + +@testset "style threads through the unfold" begin + # `one!` folds with the caller-supplied style and must unfold with the same style, not one + # re-derived from the fused matrix (here a plain `Matrix`, whose derived style would be + # `ReshapeMatricize` and would not know how to scatter into a `MyArray`). + A = MyArray(randn(3, 3)) + TA.one!(MyArrayMatricize(), A, Val(1)) + @test A.parent ≈ I +end From 6441275c24b088ca8c1e9597fda0e47d554bf82b Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Thu, 10 Sep 2026 17:13:02 -0400 Subject: [PATCH 4/5] Loosen the degenerate-truncation test tolerance above BLAS noise Co-Authored-By: Claude Fable 5 --- test/test_matrixalgebra.jl | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/test/test_matrixalgebra.jl b/test/test_matrixalgebra.jl index 291db7a8..e279ea4d 100644 --- a/test/test_matrixalgebra.jl +++ b/test/test_matrixalgebra.jl @@ -21,10 +21,13 @@ elts = (Float32, Float64, ComplexF32, ComplexF64) @test size(ṽ) == (n, n) @test ũ * s̃ * ṽ ≈ a + # The absolute tolerance must sit above the BLAS-dependent noise in the recovered + # degenerate pair (about 1eps on some builds) while staying far below the smallest + # genuine gap in the spectrum. for kwargs in ( - (; atol = eps(real(elt))), + (; atol = 100eps(real(elt))), (; rtol = (√eps(real(elt)))), - (; atol = eps(real(elt)), rtol = (√eps(real(elt)))), + (; atol = 100eps(real(elt)), rtol = (√eps(real(elt)))), ) ũ, s̃, ṽ = svd_trunc(a; trunc = truncdegen(truncrank(5); kwargs...)) @test size(ũ) == (n, 4) From 4b581bdc3ffc1b51bf465d16b824fc75f2924035 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Thu, 10 Sep 2026 18:35:38 -0400 Subject: [PATCH 5/5] Make unmatricize_factors unfold the factors only Co-Authored-By: Claude Fable 5 --- src/factorizations.jl | 89 ++++++++++++++++++++----------------------- 1 file changed, 41 insertions(+), 48 deletions(-) diff --git a/src/factorizations.jl b/src/factorizations.jl index 7a98828e..5f108cf3 100644 --- a/src/factorizations.jl +++ b/src/factorizations.jl @@ -9,16 +9,15 @@ using MatrixAlgebraKit: MatrixAlgebraKit # bond is dualized to codomain-facing form (`conj`, a no-op on a dense axis) when it lands on the # domain side of the reconstruction, matching the `unmatricize`/`similar_map` axis convention. -# `unmatricize_factors(f, style, A_mat, axes_codomain, axes_domain; kwargs...)` is the shared -# matrix-level body of the wrapper `f`: apply the matrix-level function to the matricized input -# `A_mat` and unfold the outputs with the bipartitioned axes (in the `unmatricize` convention, -# domain axes un-dualized). +# `unmatricize_factors(f, style, F, axes_codomain, axes_domain)` unfolds the matrix-level +# factors `F` of `f` onto the split axes (in the `unmatricize` convention, domain axes +# un-dualized). # # Owned tier: the matrix-level entries mutate their input, so the perm form materializes an # owned matricization following MatrixAlgebraKit's `f(A) = f!(copy_input(f, A))` convention — # a memory-sharing matricization is materialized through `MatrixAlgebraKit.copy_input`, while # the `matricizecopy` gather is owned by contract and is donated directly (with -# `copy_input` still applied when the eltype must change) — and the cores call the mutating +# `copy_input` still applied when the eltype must change) — and the wrapper calls the mutating # entry unconditionally. for f in ( :qr_compact, :qr_full, :lq_compact, :lq_full, @@ -46,13 +45,12 @@ for f in ( MatrixAlgebraKit.copy_input(MatrixAlgebraKit.$f, A_gather) end end + F = MatrixAlgebraKit.$(Symbol(f, :!))(A_mat; kwargs...) axes_codomain, axes_domain = bipartition_axes( map(i -> axes(A, i), (perm_codomain..., perm_domain...)), Val(length(perm_codomain)) ) - return unmatricize_factors( - $f, style, A_mat, axes_codomain, axes_domain; kwargs... - ) + return unmatricize_factors($f, style, F, axes_codomain, axes_domain) end end end @@ -70,13 +68,12 @@ for f in ( kwargs... ) A_mat = matricizeperm(style, A, perm_codomain, perm_domain) + F = MatrixAlgebra.$f(A_mat; kwargs...) axes_codomain, axes_domain = bipartition_axes( map(i -> axes(A, i), (perm_codomain..., perm_domain...)), Val(length(perm_codomain)) ) - return unmatricize_factors( - $f, style, A_mat, axes_codomain, axes_domain; kwargs... - ) + return unmatricize_factors($f, style, F, axes_codomain, axes_domain) end end end @@ -134,10 +131,10 @@ for f in ( ) @eval begin function unmatricize_factors( - ::typeof($f), style::MatricizeStyle, A_mat, - axes_codomain, axes_domain; kwargs... + ::typeof($f), style::MatricizeStyle, F, + axes_codomain, axes_domain ) - X, Y = MatrixAlgebraKit.$(Symbol(f, :!))(A_mat; kwargs...) + X, Y = F return unmatricize(style, X, axes_codomain, (conj(axes(X, ndims(X))),)), unmatricize(style, Y, (axes(Y, 1),), axes_domain) end @@ -330,10 +327,10 @@ right_orth for f in (:svd_compact, :svd_full) @eval begin function unmatricize_factors( - ::typeof($f), style::MatricizeStyle, A_mat, - axes_codomain, axes_domain; kwargs... + ::typeof($f), style::MatricizeStyle, F, + axes_codomain, axes_domain ) - U, S, Vᴴ = MatrixAlgebraKit.$(Symbol(f, :!))(A_mat; kwargs...) + U, S, Vᴴ = F return unmatricize(style, U, axes_codomain, (conj(axes(U, ndims(U))),)), S, unmatricize(style, Vᴴ, (axes(Vᴴ, 1),), axes_domain) @@ -345,10 +342,10 @@ end # `ϵ` (the 2-norm of the discarded singular values, computed by MatrixAlgebraKit without # catastrophic cancellation), so it is spelled out here rather than sharing the loop above. function unmatricize_factors( - ::typeof(svd_trunc), style::MatricizeStyle, A_mat, - axes_codomain, axes_domain; kwargs... + ::typeof(svd_trunc), style::MatricizeStyle, F, + axes_codomain, axes_domain ) - U, S, Vᴴ, ϵ = MatrixAlgebraKit.svd_trunc!(A_mat; kwargs...) + U, S, Vᴴ, ϵ = F return unmatricize(style, U, axes_codomain, (conj(axes(U, ndims(U))),)), S, unmatricize(style, Vᴴ, (axes(Vᴴ, 1),), axes_domain), @@ -361,24 +358,25 @@ end for f in (:eigh_full, :eig_full, :eigh_trunc, :eig_trunc) @eval begin function unmatricize_factors( - ::typeof($f), style::MatricizeStyle, A_mat, - axes_codomain, axes_domain; kwargs... + ::typeof($f), style::MatricizeStyle, F, + axes_codomain, axes_domain ) - D, V = MatrixAlgebraKit.$(Symbol(f, :!))(A_mat; kwargs...) + D, V = F return D, unmatricize(style, V, axes_codomain, (conj(axes(V, ndims(V))),)) end end end -# Spectrum-only factorizations returning a vector of singular values / eigenvalues. +# Spectrum-only factorizations returning a vector of singular values / eigenvalues: +# nothing to unfold. for f in (:svd_vals, :eigh_vals, :eig_vals) @eval begin function unmatricize_factors( - ::typeof($f), style::MatricizeStyle, A_mat, - axes_codomain, axes_domain; kwargs... + ::typeof($f), ::MatricizeStyle, F, + axes_codomain, axes_domain ) - return MatrixAlgebraKit.$(Symbol(f, :!))(A_mat; kwargs...) + return F end end end @@ -566,10 +564,9 @@ function left_null!!(A, ndims_codomain::Val; kwargs...) end function unmatricize_factors( - ::typeof(left_null), style::MatricizeStyle, A_mat, - axes_codomain, axes_domain; kwargs... + ::typeof(left_null), style::MatricizeStyle, N, + axes_codomain, axes_domain ) - N = MatrixAlgebraKit.left_null!(A_mat; kwargs...) return unmatricize(style, N, axes_codomain, (conj(axes(N, ndims(N))),)) end @@ -604,10 +601,9 @@ function right_null!!(A, ndims_codomain::Val; kwargs...) end function unmatricize_factors( - ::typeof(right_null), style::MatricizeStyle, A_mat, - axes_codomain, axes_domain; kwargs... + ::typeof(right_null), style::MatricizeStyle, Nᴴ, + axes_codomain, axes_domain ) - Nᴴ = MatrixAlgebraKit.right_null!(A_mat; kwargs...) return unmatricize(style, Nᴴ, (axes(Nᴴ, 1),), axes_domain) end @@ -661,10 +657,9 @@ function gram_eigh_full!!(A, ndims_codomain::Val; kwargs...) end function unmatricize_factors( - ::typeof(gram_eigh_full), style::MatricizeStyle, A_mat, - axes_codomain, axes_domain; kwargs... + ::typeof(gram_eigh_full), style::MatricizeStyle, X, + axes_codomain, axes_domain ) - X = MatrixAlgebra.gram_eigh_full(A_mat; kwargs...) return unmatricize(style, X, axes_codomain, (conj(axes(X, ndims(X))),)) end @@ -722,10 +717,10 @@ function gram_eigh_full_with_pinv!!(A, ndims_codomain::Val; kwargs...) end function unmatricize_factors( - ::typeof(gram_eigh_full_with_pinv), style::MatricizeStyle, A_mat, - axes_codomain, axes_domain; kwargs... + ::typeof(gram_eigh_full_with_pinv), style::MatricizeStyle, F, + axes_codomain, axes_domain ) - X, Y = MatrixAlgebra.gram_eigh_full_with_pinv(A_mat; kwargs...) + X, Y = F return unmatricize(style, X, axes_codomain, (conj(axes(X, ndims(X))),)), unmatricize(style, Y, (axes(Y, 1),), axes_codomain) end @@ -779,10 +774,9 @@ invsqrth_safe for f in (:sqrth_safe, :invsqrth_safe) @eval begin function unmatricize_factors( - ::typeof($f), style::MatricizeStyle, A_mat, - axes_codomain, axes_domain; kwargs... + ::typeof($f), style::MatricizeStyle, P_mat, + axes_codomain, axes_domain ) - P_mat = MatrixAlgebra.$f(A_mat; kwargs...) return unmatricize(style, P_mat, axes_codomain, axes_domain) end end @@ -802,10 +796,9 @@ See also `MatrixAlgebraKit.project_hermitian`. project_hermitian function unmatricize_factors( - ::typeof(project_hermitian), style::MatricizeStyle, A_mat, - axes_codomain, axes_domain; kwargs... + ::typeof(project_hermitian), style::MatricizeStyle, H_mat, + axes_codomain, axes_domain ) - H_mat = MatrixAlgebraKit.project_hermitian!(A_mat; kwargs...) return unmatricize(style, H_mat, axes_codomain, axes_domain) end @@ -830,10 +823,10 @@ See also [`MatrixAlgebra.sqrth_invsqrth_safe`](@ref). sqrth_invsqrth_safe function unmatricize_factors( - ::typeof(sqrth_invsqrth_safe), style::MatricizeStyle, A_mat, - axes_codomain, axes_domain; kwargs... + ::typeof(sqrth_invsqrth_safe), style::MatricizeStyle, F, + axes_codomain, axes_domain ) - P_mat, Pinv_mat = MatrixAlgebra.sqrth_invsqrth_safe(A_mat; kwargs...) + P_mat, Pinv_mat = F return unmatricize(style, P_mat, axes_codomain, axes_domain), unmatricize(style, Pinv_mat, axes_codomain, axes_domain) end