From afe0300bfba4daf72224b6107b0ead92fbe13f51 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Wed, 1 Jul 2026 14:50:39 -0400 Subject: [PATCH] Loosen generic tensor operations to accept non-array operands Loosens the generic tensor operations (contraction, factorizations, matrix functions, permuted-add) to accept operands that are not `AbstractArray`s, so a backend can flow its own tensor type through the generic orchestration by overloading a small set of hooks rather than subtyping `AbstractArray`. The supporting refactors: - `matricize`/`unmatricize` split into permuting and non-permuting forms by name: the permuting forms are `matricizeperm`/`unmatricizeperm`/`unmatricizeperm!` (and `matricizeop` becomes `matricizeopperm`), while `unmatricize` is again the split-axes primitive `unmatricize(m, codomain_axes, domain_axes)`. Distinguishing the two resolves the method ambiguities that had forced the flattened signature and restores the split-axes call. - Output allocation goes through the bipartitioned `similar_map` hook instead of a flat `similar`, so a map-shaped backend can build its codomain and domain directly. - A private `PermutedDims` lazy-permutation wrapper mirrors `PermutedDimsArray` for non-array operands, letting a downstream `permuteddims` hand a lazily-permuted operand to `bipermutedimsopadd!`. The public `matricizeop` to `matricizeopperm` rename is breaking. The TensorKit backend that exercises these hooks lands in a follow-up PR. --- Project.toml | 2 +- docs/Project.toml | 2 +- examples/Project.toml | 2 +- .../TensorAlgebraTensorOperationsExt.jl | 6 +- src/TensorAlgebra.jl | 2 +- src/contract/allocate_output.jl | 27 +-- src/contract/contract.jl | 60 ++--- src/contract/contract_matricize.jl | 14 +- src/contract/contractalgorithm.jl | 12 +- src/factorizations.jl | 205 +++++++++--------- src/linearbroadcasted.jl | 15 +- src/matricize.jl | 142 ++++++------ src/matrixfunctions.jl | 15 +- src/permutedimsadd.jl | 42 +++- src/similar_map.jl | 11 +- test/Project.toml | 2 +- test/test_basics.jl | 68 +++--- test/test_exports.jl | 2 +- test/test_matricize.jl | 20 +- test/test_permutedimsadd.jl | 40 +++- 20 files changed, 371 insertions(+), 318 deletions(-) diff --git a/Project.toml b/Project.toml index bb499329..31e1b697 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "TensorAlgebra" uuid = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a" -version = "0.14.1" +version = "0.15.0" authors = ["ITensor developers and contributors"] [workspace] diff --git a/docs/Project.toml b/docs/Project.toml index 6bb4067a..119db360 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -11,4 +11,4 @@ path = ".." Documenter = "1.8.1" ITensorFormatter = "0.2.27" Literate = "2.20.1" -TensorAlgebra = "0.14" +TensorAlgebra = "0.15" diff --git a/examples/Project.toml b/examples/Project.toml index 96de614b..9139910e 100644 --- a/examples/Project.toml +++ b/examples/Project.toml @@ -5,4 +5,4 @@ TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a" path = ".." [compat] -TensorAlgebra = "0.14" +TensorAlgebra = "0.15" diff --git a/ext/TensorAlgebraTensorOperationsExt/TensorAlgebraTensorOperationsExt.jl b/ext/TensorAlgebraTensorOperationsExt/TensorAlgebraTensorOperationsExt.jl index 778b7f58..7cba5cac 100644 --- a/ext/TensorAlgebraTensorOperationsExt/TensorAlgebraTensorOperationsExt.jl +++ b/ext/TensorAlgebraTensorOperationsExt/TensorAlgebraTensorOperationsExt.jl @@ -54,9 +54,9 @@ end # in-place function TA.contractopadd!( algorithm::TensorOperationsAlgorithm, - a_dest::AbstractArray, perm_dest_codomain, perm_dest_domain, - op1, a1::AbstractArray, perm1_codomain, perm1_domain, - op2, a2::AbstractArray, perm2_codomain, perm2_domain, + a_dest, perm_dest_codomain, perm_dest_domain, + op1, a1, perm1_codomain, perm1_domain, + op2, a2, perm2_codomain, perm2_domain, α::Number, β::Number ) permblocks1 = Tuple.((perm1_codomain, perm1_domain)) diff --git a/src/TensorAlgebra.jl b/src/TensorAlgebra.jl index 55b3fdbf..f9d5366d 100644 --- a/src/TensorAlgebra.jl +++ b/src/TensorAlgebra.jl @@ -8,7 +8,7 @@ export contract, contract!, eig_full, eig_trunc, eig_vals, eigh_full, eigh_trunc if VERSION >= v"1.11.0-DEV.469" eval( Meta.parse( - "public biperm, bipartition, contractopadd!, label_type, matricizeop, to_range, zero!, scale!, permuteddims, conjed, ConjArray" + "public biperm, bipartition, contractopadd!, label_type, matricizeopperm, to_range, zero!, scale!, permuteddims, conjed, ConjArray" ) ) end diff --git a/src/contract/allocate_output.jl b/src/contract/allocate_output.jl index 2f71f6c5..cebee562 100644 --- a/src/contract/allocate_output.jl +++ b/src/contract/allocate_output.jl @@ -1,5 +1,3 @@ -using Base.PermutedDimsArrays: genperm - function check_biperm(a, perm_codomain, perm_domain) ndims(a) == length(perm_codomain) + length(perm_domain) || throw(ArgumentError("Invalid bipartitioned permutation")) @@ -44,16 +42,13 @@ end function output_axes( ::typeof(contract), perm_dest_codomain, perm_dest_domain, - a1::AbstractArray, perm1_codomain, perm1_domain, - a2::AbstractArray, perm2_codomain, perm2_domain - ) - axes_codomain, axes_contracted = bipartition(axes(a1), perm1_codomain, perm1_domain) - axes_contracted2, axes_domain = bipartition(axes(a2), perm2_codomain, perm2_domain) - @assert length.(axes_contracted) == length.(axes_contracted2) - # default: flatten the destination permutation - return genperm( - (axes_codomain..., axes_domain...), (perm_dest_codomain..., perm_dest_domain...) + a1, perm1_codomain, perm1_domain, + a2, perm2_codomain, perm2_domain ) + axes_codomain, _ = bipartition(axes(a1), perm1_codomain, perm1_domain) + _, axes_domain = bipartition(axes(a2), perm2_codomain, perm2_domain) + axes_uncontracted = (axes_codomain..., axes_domain...) + return bipartition(axes_uncontracted, perm_dest_codomain, perm_dest_domain) end # TODO: Use `ArrayLayouts`-like `MulAdd` object, @@ -61,8 +56,8 @@ end function allocate_output( ::typeof(contract), perm_dest_codomain, perm_dest_domain, - a1::AbstractArray, perm1_codomain, perm1_domain, - a2::AbstractArray, perm2_codomain, perm2_domain + a1, perm1_codomain, perm1_domain, + a2, perm2_codomain, perm2_domain ) check_input( contract, @@ -73,12 +68,14 @@ function allocate_output( perm2_codomain, perm2_domain ) - axes_dest = output_axes( + codomain_axes_dest, domain_axes_dest = output_axes( contract, perm_dest_codomain, perm_dest_domain, a1, perm1_codomain, perm1_domain, a2, perm2_codomain, perm2_domain ) T = promote_type(eltype(a1), eltype(a2)) - return zero!(similar(a1, T, axes_dest)) + # `domain_axes_dest` come straight from `axes(a2)` (stored/dualized convention), so + # un-dualize them into `similar_map`'s codomain-facing convention. + return zero!(similar_map(a1, T, codomain_axes_dest, conj.(domain_axes_dest))) end diff --git a/src/contract/contract.jl b/src/contract/contract.jl index 426936ed..f31f14ea 100644 --- a/src/contract/contract.jl +++ b/src/contract/contract.jl @@ -2,7 +2,7 @@ # TODO: Add `scaledcontract(a1, labels1, a2, labels2, α) = α * contract(a1, labels1, a2, labels2)`. # contract (labels) -function contract(a1::AbstractArray, labels1, a2::AbstractArray, labels2; kwargs...) +function contract(a1, labels1, a2, labels2; kwargs...) # Optionally convert the labels to a representation cheaper to run the bookkeeping on (see # `label_type`). `encode_contraction_labels`/`decode_contraction_labels` are no-ops unless the label type opts in. l1, l2 = encode_contraction_labels(labels1, labels2) @@ -11,7 +11,7 @@ function contract(a1::AbstractArray, labels1, a2::AbstractArray, labels2; kwargs return a_dest, decode_contraction_labels(l_dest, labels1, labels2) end function contract( - labels_dest, a1::AbstractArray, labels1, a2::AbstractArray, labels2; kwargs... + labels_dest, a1, labels1, a2, labels2; kwargs... ) t1 = ntuple(i -> labels1[i], Val(ndims(a1))) t2 = ntuple(i -> labels2[i], Val(ndims(a2))) @@ -30,7 +30,7 @@ function contract( ) end function _contract( - ::Val{K}, labels_dest, a1::AbstractArray, labels1, a2::AbstractArray, labels2, + ::Val{K}, labels_dest, a1, labels1, a2, labels2, contracted1; kwargs... ) where {K} biperm_dest, biperm1, biperm2 = @@ -40,8 +40,8 @@ end # contract (bipartitioned permutations) function contract( - a1::AbstractArray, perm1_codomain, perm1_domain, - a2::AbstractArray, perm2_codomain, perm2_domain; + a1, perm1_codomain, perm1_domain, + a2, perm2_codomain, perm2_domain; kwargs... ) Ndest_codomain = Val(length(perm1_codomain)) @@ -57,8 +57,8 @@ function contract( end function contract( perm_dest_codomain, perm_dest_domain, - a1::AbstractArray, perm1_codomain, perm1_domain, - a2::AbstractArray, perm2_codomain, perm2_domain; + a1, perm1_codomain, perm1_domain, + a2, perm2_codomain, perm2_domain; kwargs... ) a_dest = allocate_output( @@ -77,9 +77,9 @@ end # contract! (labels) function contract!( - a_dest::AbstractArray, labels_dest, - a1::AbstractArray, labels1, - a2::AbstractArray, labels2; + a_dest, labels_dest, + a1, labels1, + a2, labels2; kwargs... ) return contractadd!( @@ -87,9 +87,9 @@ function contract!( ) end function contract!( - a_dest::AbstractArray, perm_dest_codomain, perm_dest_domain, - a1::AbstractArray, perm1_codomain, perm1_domain, - a2::AbstractArray, perm2_codomain, perm2_domain; + a_dest, perm_dest_codomain, perm_dest_domain, + a1, perm1_codomain, perm1_domain, + a2, perm2_codomain, perm2_domain; kwargs... ) return contractadd!( @@ -102,9 +102,9 @@ end # contractadd! (labels) function contractadd!( - a_dest::AbstractArray, labels_dest, - a1::AbstractArray, labels1, - a2::AbstractArray, labels2, + a_dest, labels_dest, + a1, labels1, + a2, labels2, α::Number, β::Number; kwargs... ) @@ -114,9 +114,9 @@ function contractadd!( end # contractadd! (bipartitioned permutations) function contractadd!( - a_dest::AbstractArray, perm_dest_codomain, perm_dest_domain, - a1::AbstractArray, perm1_codomain, perm1_domain, - a2::AbstractArray, perm2_codomain, perm2_domain, + a_dest, perm_dest_codomain, perm_dest_domain, + a1, perm1_codomain, perm1_domain, + a2, perm2_codomain, perm2_domain, α::Number, β::Number; kwargs... ) @@ -130,9 +130,9 @@ end # contractopadd! (labels) function contractopadd!( - a_dest::AbstractArray, labels_dest, - op1, a1::AbstractArray, labels1, - op2, a2::AbstractArray, labels2, + a_dest, labels_dest, + op1, a1, labels1, + op2, a2, labels2, α::Number, β::Number; kwargs... ) @@ -147,8 +147,8 @@ function contractopadd!( ) end function _contractopadd!( - ::Val{K}, a_dest::AbstractArray, labels_dest, - op1, a1::AbstractArray, labels1, op2, a2::AbstractArray, labels2, + ::Val{K}, a_dest, labels_dest, + op1, a1, labels1, op2, a2, labels2, α::Number, β::Number, contracted1; kwargs... ) where {K} biperm_dest, biperm1, biperm2 = @@ -159,9 +159,9 @@ function _contractopadd!( end # contractopadd! (bipartitioned permutations, algorithm selection) function contractopadd!( - a_dest::AbstractArray, perm_dest_codomain, perm_dest_domain, - op1, a1::AbstractArray, perm1_codomain, perm1_domain, - op2, a2::AbstractArray, perm2_codomain, perm2_domain, + a_dest, perm_dest_codomain, perm_dest_domain, + op1, a1, perm1_codomain, perm1_domain, + op2, a2, perm2_codomain, perm2_domain, α::Number, β::Number; alg = DefaultContractAlgorithm(), kwargs... ) @@ -184,9 +184,9 @@ end # Required interface if not using matricized contraction function contractopadd!( algorithm::ContractAlgorithm, - a_dest::AbstractArray, perm_dest_codomain, perm_dest_domain, - op1, a1::AbstractArray, perm1_codomain, perm1_domain, - op2, a2::AbstractArray, perm2_codomain, perm2_domain, + a_dest, perm_dest_codomain, perm_dest_domain, + op1, a1, perm1_codomain, perm1_domain, + op2, a2, perm2_codomain, perm2_domain, α::Number, β::Number ) return throw( diff --git a/src/contract/contract_matricize.jl b/src/contract/contract_matricize.jl index 75253998..b83277d6 100644 --- a/src/contract/contract_matricize.jl +++ b/src/contract/contract_matricize.jl @@ -16,20 +16,22 @@ function contractopadd!( a1, biperm1_codomain, biperm1_domain, a2, biperm2_codomain, biperm2_domain ) - a1_mat = - matricizeop(algorithm.left_fusion_style, op1, a1, biperm1_codomain, biperm1_domain) - a2_mat = - matricizeop(algorithm.right_fusion_style, op2, a2, biperm2_codomain, biperm2_domain) + a1_mat = matricizeopperm( + algorithm.left_fusion_style, op1, a1, biperm1_codomain, biperm1_domain + ) + a2_mat = matricizeopperm( + algorithm.right_fusion_style, op2, a2, biperm2_codomain, biperm2_domain + ) output_style = algorithm.output_fusion_style # 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_mat = matricize(output_style, a_dest, invperm_codomain, invperm_domain) + 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) - unmatricize!(output_style, a_dest, a_dest_mat, invperm_codomain, invperm_domain) + unmatricizeperm!(output_style, a_dest, a_dest_mat, invperm_codomain, invperm_domain) end return a_dest end diff --git a/src/contract/contractalgorithm.jl b/src/contract/contractalgorithm.jl index 48bda4f7..46c929ee 100644 --- a/src/contract/contractalgorithm.jl +++ b/src/contract/contractalgorithm.jl @@ -11,20 +11,16 @@ end Matricize(fusion_style) = Matricize(fusion_style, fusion_style, fusion_style) Matricize() = Matricize(ReshapeFusion()) -function select_contract_algorithm(algorithm, a1::AbstractArray, a2::AbstractArray) +function select_contract_algorithm(algorithm, a1, a2) return error("Not implemented.") end -function select_contract_algorithm( - algorithm::ContractAlgorithm, a1::AbstractArray, a2::AbstractArray - ) +function select_contract_algorithm(algorithm::ContractAlgorithm, a1, a2) return algorithm end -function select_contract_algorithm( - algorithm::DefaultContractAlgorithm, a1::AbstractArray, a2::AbstractArray - ) +function select_contract_algorithm(algorithm::DefaultContractAlgorithm, a1, a2) return default_contract_algorithm(a1, a2) end -function default_contract_algorithm(a1::AbstractArray, a2::AbstractArray) +function default_contract_algorithm(a1, a2) return default_contract_algorithm(typeof(a1), typeof(a2)) end function default_contract_algorithm(A1::Type{<:AbstractArray}, A2::Type{<:AbstractArray}) diff --git a/src/factorizations.jl b/src/factorizations.jl index 3d5ff050..c4ab7a9e 100644 --- a/src/factorizations.jl +++ b/src/factorizations.jl @@ -8,14 +8,14 @@ for f in ( :left_polar, :right_polar, :left_orth, :right_orth, ) @eval begin - function $f(style::FusionStyle, A::AbstractArray, ndims_codomain::Val; kwargs...) + function $f(style::FusionStyle, 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(A), ndims_codomain) return unmatricize(style, X, axes_codomain, (axes(X, 2),)), unmatricize(style, Y, (axes(Y, 1),), axes_domain) end - function $f(A::AbstractArray, ndims_codomain::Val; kwargs...) + function $f(A, ndims_codomain::Val; kwargs...) return $f(FusionStyle(A), A, ndims_codomain; kwargs...) end end @@ -30,7 +30,7 @@ for f in ( ) @eval begin function $f( - style::FusionStyle, A::AbstractArray, + style::FusionStyle, A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs... ) @@ -38,7 +38,7 @@ for f in ( return $f(style, A_perm, Val(length(perm_codomain)); kwargs...) end function $f( - A::AbstractArray, + A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs... ) @@ -47,14 +47,14 @@ for f in ( end function $f( - style::FusionStyle, A::AbstractArray, + style::FusionStyle, A, labels_A, labels_codomain, labels_domain; kwargs... ) perm_codomain, perm_domain = biperm(Tuple.((labels_A, labels_codomain, labels_domain))...) return $f(style, A, perm_codomain, perm_domain; kwargs...) end - function $f(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) + function $f(A, labels_A, labels_codomain, labels_domain; kwargs...) perm_codomain, perm_domain = biperm(Tuple.((labels_A, labels_codomain, labels_domain))...) return $f(A, perm_codomain, perm_domain; kwargs...) @@ -63,9 +63,9 @@ for f in ( end """ - qr_compact(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> Q, R - qr_compact(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> Q, R - qr_compact(A::AbstractArray, ndims_codomain::Val; kwargs...) -> Q, R + qr_compact(A, labels_A, labels_codomain, labels_domain; kwargs...) -> Q, R + qr_compact(A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> Q, R + qr_compact(A, ndims_codomain::Val; kwargs...) -> Q, R Compute the compact QR decomposition of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions, where `R` is square. The @@ -81,9 +81,9 @@ See also `MatrixAlgebraKit.qr_compact!`. qr_compact """ - qr_full(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> Q, R - qr_full(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> Q, R - qr_full(A::AbstractArray, ndims_codomain::Val; kwargs...) -> Q, R + qr_full(A, labels_A, labels_codomain, labels_domain; kwargs...) -> Q, R + qr_full(A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> Q, R + qr_full(A, ndims_codomain::Val; kwargs...) -> Q, R Compute the full QR decomposition of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions, where `Q` is unitary. The @@ -99,9 +99,9 @@ See also `MatrixAlgebraKit.qr_full!`. qr_full """ - lq_compact(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> L, Q - lq_compact(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> L, Q - lq_compact(A::AbstractArray, ndims_codomain::Val; kwargs...) -> L, Q + lq_compact(A, labels_A, labels_codomain, labels_domain; kwargs...) -> L, Q + lq_compact(A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> L, Q + lq_compact(A, ndims_codomain::Val; kwargs...) -> L, Q Compute the compact LQ decomposition of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions, where `L` is square. The @@ -117,9 +117,9 @@ See also `MatrixAlgebraKit.lq_compact!`. lq_compact """ - lq_full(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> L, Q - lq_full(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> L, Q - lq_full(A::AbstractArray, ndims_codomain::Val; kwargs...) -> L, Q + lq_full(A, labels_A, labels_codomain, labels_domain; kwargs...) -> L, Q + lq_full(A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> L, Q + lq_full(A, ndims_codomain::Val; kwargs...) -> L, Q Compute the full LQ decomposition of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions, where `Q` is unitary. The @@ -135,9 +135,9 @@ See also `MatrixAlgebraKit.lq_full!`. lq_full """ - left_polar(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> W, P - left_polar(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> W, P - left_polar(A::AbstractArray, ndims_codomain::Val; kwargs...) -> W, P + left_polar(A, labels_A, labels_codomain, labels_domain; kwargs...) -> W, P + left_polar(A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> W, P + left_polar(A, ndims_codomain::Val; kwargs...) -> W, P Compute the left polar decomposition of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions. These can be specified either via @@ -152,9 +152,9 @@ See also `MatrixAlgebraKit.left_polar!`. left_polar """ - right_polar(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> P, W - right_polar(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> P, W - right_polar(A::AbstractArray, ndims_codomain::Val; kwargs...) -> P, W + right_polar(A, labels_A, labels_codomain, labels_domain; kwargs...) -> P, W + right_polar(A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> P, W + right_polar(A, ndims_codomain::Val; kwargs...) -> P, W Compute the right polar decomposition of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions. These can be specified either via @@ -169,9 +169,9 @@ See also `MatrixAlgebraKit.right_polar!`. right_polar """ - left_orth(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> V, C - left_orth(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> V, C - left_orth(A::AbstractArray, ndims_codomain::Val; kwargs...) -> V, C + left_orth(A, labels_A, labels_codomain, labels_domain; kwargs...) -> V, C + left_orth(A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> V, C + left_orth(A, ndims_codomain::Val; kwargs...) -> V, C Compute the left orthogonal decomposition of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions. These can be specified either via @@ -186,9 +186,9 @@ See also `MatrixAlgebraKit.left_orth!`. left_orth """ - right_orth(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> C, V - right_orth(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> C, V - right_orth(A::AbstractArray, ndims_codomain::Val; kwargs...) -> C, V + right_orth(A, labels_A, labels_codomain, labels_domain; kwargs...) -> C, V + right_orth(A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> C, V + right_orth(A, ndims_codomain::Val; kwargs...) -> C, V Compute the right orthogonal decomposition of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions. These can be specified either via @@ -206,7 +206,7 @@ right_orth # rank × rank spectrum, and `Vᴴ` carries a leading rank axis plus the domain axes. for f in (:svd_compact, :svd_full, :svd_trunc) @eval begin - function $f(style::FusionStyle, A::AbstractArray, ndims_codomain::Val; kwargs...) + function $f(style::FusionStyle, 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(A), ndims_codomain) @@ -214,7 +214,7 @@ for f in (:svd_compact, :svd_full, :svd_trunc) unmatricize(style, S, (axes(S, 1),), (axes(S, 2),)), unmatricize(style, Vᴴ, (axes(Vᴴ, 1),), axes_domain) end - function $f(A::AbstractArray, ndims_codomain::Val; kwargs...) + function $f(A, ndims_codomain::Val; kwargs...) return $f(FusionStyle(A), A, ndims_codomain; kwargs...) end end @@ -224,13 +224,13 @@ end # the codomain axes plus a trailing rank axis. for f in (:eigh_full, :eig_full, :eigh_trunc, :eig_trunc) @eval begin - function $f(style::FusionStyle, A::AbstractArray, ndims_codomain::Val; kwargs...) + function $f(style::FusionStyle, 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)) return D, unmatricize(style, V, axes_codomain, (axes(V, ndims(V)),)) end - function $f(A::AbstractArray, ndims_codomain::Val; kwargs...) + function $f(A, ndims_codomain::Val; kwargs...) return $f(FusionStyle(A), A, ndims_codomain; kwargs...) end end @@ -239,20 +239,20 @@ 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::FusionStyle, A::AbstractArray, ndims_codomain::Val; kwargs...) + function $f(style::FusionStyle, A, ndims_codomain::Val; kwargs...) A_mat = matricize(style, A, ndims_codomain) return MatrixAlgebraKit.$f(A_mat; kwargs...) end - function $f(A::AbstractArray, ndims_codomain::Val; kwargs...) + function $f(A, ndims_codomain::Val; kwargs...) return $f(FusionStyle(A), A, ndims_codomain; kwargs...) end end end """ - svd_compact(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> U, S, Vᴴ - svd_compact(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> U, S, Vᴴ - svd_compact(A::AbstractArray, ndims_codomain::Val; kwargs...) -> U, S, Vᴴ + svd_compact(A, labels_A, labels_codomain, labels_domain; kwargs...) -> U, S, Vᴴ + svd_compact(A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> U, S, Vᴴ + svd_compact(A, ndims_codomain::Val; kwargs...) -> U, S, Vᴴ Compute the compact (thin) SVD of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions, where `U` and `Vᴴ` are isometric. @@ -263,9 +263,9 @@ See also `MatrixAlgebraKit.svd_compact!`. svd_compact """ - svd_full(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> U, S, Vᴴ - svd_full(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> U, S, Vᴴ - svd_full(A::AbstractArray, ndims_codomain::Val; kwargs...) -> U, S, Vᴴ + svd_full(A, labels_A, labels_codomain, labels_domain; kwargs...) -> U, S, Vᴴ + svd_full(A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> U, S, Vᴴ + svd_full(A, ndims_codomain::Val; kwargs...) -> U, S, Vᴴ Compute the full (thick) SVD of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions, where `U` and `Vᴴ` are unitary. @@ -276,9 +276,9 @@ See also `MatrixAlgebraKit.svd_full!`. svd_full """ - svd_trunc(A::AbstractArray, labels_A, labels_codomain, labels_domain; trunc, kwargs...) -> U, S, Vᴴ - svd_trunc(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; trunc, kwargs...) -> U, S, Vᴴ - svd_trunc(A::AbstractArray, ndims_codomain::Val; trunc, kwargs...) -> U, S, Vᴴ + svd_trunc(A, labels_A, labels_codomain, labels_domain; trunc, kwargs...) -> U, S, Vᴴ + svd_trunc(A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; trunc, kwargs...) -> U, S, Vᴴ + svd_trunc(A, ndims_codomain::Val; trunc, kwargs...) -> U, S, Vᴴ Compute the truncated SVD of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions. The partition is specified either via @@ -294,9 +294,9 @@ See also `MatrixAlgebraKit.svd_trunc!`. svd_trunc """ - svd_vals(A::AbstractArray, labels_A, labels_codomain, labels_domain) -> S - svd_vals(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}) -> S - svd_vals(A::AbstractArray, ndims_codomain::Val) -> S + svd_vals(A, labels_A, labels_codomain, labels_domain) -> S + svd_vals(A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}) -> S + svd_vals(A, ndims_codomain::Val) -> S Compute the singular values of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions. The partition is specified either @@ -307,9 +307,9 @@ See also `MatrixAlgebraKit.svd_vals!`. svd_vals """ - eigh_full(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> D, V - eigh_full(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> D, V - eigh_full(A::AbstractArray, ndims_codomain::Val; kwargs...) -> D, V + eigh_full(A, labels_A, labels_codomain, labels_domain; kwargs...) -> D, V + eigh_full(A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> D, V + eigh_full(A, ndims_codomain::Val; kwargs...) -> D, V Compute the eigenvalue decomposition of a generic N-dimensional array interpreted as a Hermitian linear map from the domain to the codomain dimensions. The partition is specified @@ -320,9 +320,9 @@ See also `MatrixAlgebraKit.eigh_full!`. eigh_full """ - eig_full(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> D, V - eig_full(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> D, V - eig_full(A::AbstractArray, ndims_codomain::Val; kwargs...) -> D, V + eig_full(A, labels_A, labels_codomain, labels_domain; kwargs...) -> D, V + eig_full(A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> D, V + eig_full(A, ndims_codomain::Val; kwargs...) -> D, V Compute the eigenvalue decomposition of a generic N-dimensional array interpreted as a general (non-Hermitian) linear map from the domain to the codomain dimensions. The output @@ -334,9 +334,9 @@ See also `MatrixAlgebraKit.eig_full!`. eig_full """ - eigh_trunc(A::AbstractArray, labels_A, labels_codomain, labels_domain; trunc, kwargs...) -> D, V - eigh_trunc(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; trunc, kwargs...) -> D, V - eigh_trunc(A::AbstractArray, ndims_codomain::Val; trunc, kwargs...) -> D, V + eigh_trunc(A, labels_A, labels_codomain, labels_domain; trunc, kwargs...) -> D, V + eigh_trunc(A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; trunc, kwargs...) -> D, V + eigh_trunc(A, ndims_codomain::Val; trunc, kwargs...) -> D, V Truncated Hermitian eigenvalue decomposition, like [`eigh_full`](@ref) but keeping only the eigenvalues selected by the `trunc` strategy. @@ -346,9 +346,9 @@ See also `MatrixAlgebraKit.eigh_trunc!`. eigh_trunc """ - eig_trunc(A::AbstractArray, labels_A, labels_codomain, labels_domain; trunc, kwargs...) -> D, V - eig_trunc(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; trunc, kwargs...) -> D, V - eig_trunc(A::AbstractArray, ndims_codomain::Val; trunc, kwargs...) -> D, V + eig_trunc(A, labels_A, labels_codomain, labels_domain; trunc, kwargs...) -> D, V + eig_trunc(A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; trunc, kwargs...) -> D, V + eig_trunc(A, ndims_codomain::Val; trunc, kwargs...) -> D, V Truncated general eigenvalue decomposition, like [`eig_full`](@ref) but keeping only the eigenvalues selected by the `trunc` strategy. @@ -358,9 +358,9 @@ See also `MatrixAlgebraKit.eig_trunc!`. eig_trunc """ - eigh_vals(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> D - eigh_vals(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> D - eigh_vals(A::AbstractArray, ndims_codomain::Val; kwargs...) -> D + eigh_vals(A, labels_A, labels_codomain, labels_domain; kwargs...) -> D + eigh_vals(A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> D + eigh_vals(A, ndims_codomain::Val; kwargs...) -> D Compute the eigenvalues of a generic N-dimensional array interpreted as a Hermitian linear map from the domain to the codomain dimensions. The output is a vector of eigenvalues. @@ -370,9 +370,9 @@ See also `MatrixAlgebraKit.eigh_vals!`. eigh_vals """ - eig_vals(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> D - eig_vals(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> D - eig_vals(A::AbstractArray, ndims_codomain::Val; kwargs...) -> D + eig_vals(A, labels_A, labels_codomain, labels_domain; kwargs...) -> D + eig_vals(A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> D + eig_vals(A, ndims_codomain::Val; kwargs...) -> D Compute the eigenvalues of a generic N-dimensional array interpreted as a general (non-Hermitian) linear map from the domain to the codomain dimensions. The output is a @@ -383,9 +383,9 @@ See also `MatrixAlgebraKit.eig_vals!`. eig_vals """ - left_null(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> N - left_null(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> N - left_null(A::AbstractArray, ndims_codomain::Val; kwargs...) -> N + left_null(A, labels_A, labels_codomain, labels_domain; kwargs...) -> N + left_null(A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> N + left_null(A, ndims_codomain::Val; kwargs...) -> N Compute the left nullspace of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions. These can be specified either via @@ -402,27 +402,27 @@ The output satisfies `N' * A ≈ 0` and `N' * N ≈ I`. """ left_null -function left_null!!(style::FusionStyle, A::AbstractArray, ndims_codomain::Val; kwargs...) +function left_null!!(style::FusionStyle, A, ndims_codomain::Val; kwargs...) A_mat = matricize(style, A, ndims_codomain) N = MatrixAlgebraKit.left_null!(A_mat; kwargs...) axes_codomain = first(bipartition(axes(A), ndims_codomain)) return unmatricize(style, N, axes_codomain, (axes(N, 2),)) end -function left_null!!(A::AbstractArray, ndims_codomain::Val; kwargs...) +function left_null!!(A, ndims_codomain::Val; kwargs...) return left_null!!(FusionStyle(A), A, ndims_codomain; kwargs...) end -function left_null(style::FusionStyle, A::AbstractArray, ndims_codomain::Val; kwargs...) +function left_null(style::FusionStyle, A, ndims_codomain::Val; kwargs...) return left_null!!(style, copy(A), ndims_codomain; kwargs...) end -function left_null(A::AbstractArray, ndims_codomain::Val; kwargs...) +function left_null(A, ndims_codomain::Val; kwargs...) return left_null!!(copy(A), ndims_codomain; kwargs...) end """ - right_null(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> Nᴴ - right_null(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> Nᴴ - right_null(A::AbstractArray, ndims_codomain::Val::Val; kwargs...) -> Nᴴ + right_null(A, labels_A, labels_codomain, labels_domain; kwargs...) -> Nᴴ + right_null(A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> Nᴴ + right_null(A, ndims_codomain::Val::Val; kwargs...) -> Nᴴ Compute the right nullspace of a generic N-dimensional array, by interpreting it as a linear map from the domain to the codomain dimensions. These can be specified either via @@ -439,27 +439,27 @@ The output satisfies `A * Nᴴ' ≈ 0` and `Nᴴ * Nᴴ' ≈ I`. """ right_null -function right_null!!(style::FusionStyle, A::AbstractArray, ndims_codomain::Val; kwargs...) +function right_null!!(style::FusionStyle, A, ndims_codomain::Val; kwargs...) A_mat = matricize(style, A, ndims_codomain) Nᴴ = MatrixAlgebraKit.right_null!(A_mat; kwargs...) axes_domain = last(bipartition(axes(A), ndims_codomain)) return unmatricize(style, Nᴴ, (axes(Nᴴ, 1),), axes_domain) end -function right_null!!(A::AbstractArray, ndims_codomain::Val; kwargs...) +function right_null!!(A, ndims_codomain::Val; kwargs...) return right_null!!(FusionStyle(A), A, ndims_codomain; kwargs...) end -function right_null(style::FusionStyle, A::AbstractArray, ndims_codomain::Val; kwargs...) +function right_null(style::FusionStyle, A, ndims_codomain::Val; kwargs...) return right_null!!(style, copy(A), ndims_codomain; kwargs...) end -function right_null(A::AbstractArray, ndims_codomain::Val; kwargs...) +function right_null(A, ndims_codomain::Val; kwargs...) return right_null!!(copy(A), ndims_codomain; kwargs...) end """ - gram_eigh_full(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> X - gram_eigh_full(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> X - gram_eigh_full(A::AbstractArray, ndims_codomain::Val; kwargs...) -> X + gram_eigh_full(A, labels_A, labels_codomain, labels_domain; kwargs...) -> X + gram_eigh_full(A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> X + gram_eigh_full(A, ndims_codomain::Val; kwargs...) -> X Gram factorization of a generic N-dimensional array, interpreting it as a Hermitian positive semi-definite linear map from the domain to the codomain @@ -494,30 +494,30 @@ See also [`gram_eigh_full_with_pinv`](@ref) and gram_eigh_full function gram_eigh_full!!( - style::FusionStyle, A::AbstractArray, ndims_codomain::Val; kwargs... + style::FusionStyle, A, ndims_codomain::Val; kwargs... ) A_mat = matricize(style, A, ndims_codomain) X = MatrixAlgebra.gram_eigh_full!!(A_mat; kwargs...) axes_codomain = first(bipartition(axes(A), ndims_codomain)) return unmatricize(style, X, axes_codomain, (axes(X, 2),)) end -function gram_eigh_full!!(A::AbstractArray, ndims_codomain::Val; kwargs...) +function gram_eigh_full!!(A, ndims_codomain::Val; kwargs...) return gram_eigh_full!!(FusionStyle(A), A, ndims_codomain; kwargs...) end function gram_eigh_full( - style::FusionStyle, A::AbstractArray, ndims_codomain::Val; kwargs... + style::FusionStyle, A, ndims_codomain::Val; kwargs... ) return gram_eigh_full!!(style, copy(A), ndims_codomain; kwargs...) end -function gram_eigh_full(A::AbstractArray, ndims_codomain::Val; kwargs...) +function gram_eigh_full(A, ndims_codomain::Val; kwargs...) return gram_eigh_full!!(copy(A), ndims_codomain; kwargs...) end """ - gram_eigh_full_with_pinv(A::AbstractArray, labels_A, labels_codomain, labels_domain; kwargs...) -> X, Y - gram_eigh_full_with_pinv(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> X, Y - gram_eigh_full_with_pinv(A::AbstractArray, ndims_codomain::Val; kwargs...) -> X, Y + gram_eigh_full_with_pinv(A, labels_A, labels_codomain, labels_domain; kwargs...) -> X, Y + gram_eigh_full_with_pinv(A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs...) -> X, Y + gram_eigh_full_with_pinv(A, ndims_codomain::Val; kwargs...) -> X, Y Like [`gram_eigh_full`](@ref), but additionally returns `Y ≈ pinv(X)` such that `Y * X ≈ I` on the rank subspace (a left inverse). The codomain axes @@ -555,7 +555,7 @@ See also [`MatrixAlgebra.gram_eigh_full_with_pinv`](@ref). gram_eigh_full_with_pinv function gram_eigh_full_with_pinv!!( - style::FusionStyle, A::AbstractArray, ndims_codomain::Val; kwargs... + style::FusionStyle, A, ndims_codomain::Val; kwargs... ) A_mat = matricize(style, A, ndims_codomain) X, Y = MatrixAlgebra.gram_eigh_full_with_pinv!!(A_mat; kwargs...) @@ -563,23 +563,23 @@ function gram_eigh_full_with_pinv!!( return unmatricize(style, X, axes_codomain, (axes(X, 2),)), unmatricize(style, Y, (axes(Y, 1),), conj.(axes_codomain)) end -function gram_eigh_full_with_pinv!!(A::AbstractArray, ndims_codomain::Val; kwargs...) +function gram_eigh_full_with_pinv!!(A, ndims_codomain::Val; kwargs...) return gram_eigh_full_with_pinv!!(FusionStyle(A), A, ndims_codomain; kwargs...) end function gram_eigh_full_with_pinv( - style::FusionStyle, A::AbstractArray, ndims_codomain::Val; kwargs... + style::FusionStyle, A, ndims_codomain::Val; kwargs... ) return gram_eigh_full_with_pinv!!(style, copy(A), ndims_codomain; kwargs...) end -function gram_eigh_full_with_pinv(A::AbstractArray, ndims_codomain::Val; kwargs...) +function gram_eigh_full_with_pinv(A, ndims_codomain::Val; kwargs...) return gram_eigh_full_with_pinv!!(copy(A), ndims_codomain; kwargs...) end """ - TensorAlgebra.one(A::AbstractArray, labels_A, labels_codomain, labels_domain) -> Id - TensorAlgebra.one(A::AbstractArray, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}) -> Id - TensorAlgebra.one(A::AbstractArray, ndims_codomain::Val) -> Id + TensorAlgebra.one(A, labels_A, labels_codomain, labels_domain) -> Id + TensorAlgebra.one(A, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}) -> Id + TensorAlgebra.one(A, ndims_codomain::Val) -> Id Construct the identity operator tensor whose shape mirrors `A`, interpreted as a linear map from the domain to the codomain dimensions. The codomain and domain @@ -609,18 +609,19 @@ true """ one -function one!!(style::FusionStyle, A::AbstractArray, ndims_codomain::Val; kwargs...) +function one!!(style::FusionStyle, A, ndims_codomain::Val; kwargs...) A_mat = matricize(style, A, ndims_codomain) MatrixAlgebraKit.one!(A_mat) - return unmatricize(style, A_mat, bipartition(axes(A), ndims_codomain)...) + codomain_axes, domain_axes = bipartition(axes(A), ndims_codomain) + return unmatricize(style, A_mat, codomain_axes, domain_axes) end -function one!!(A::AbstractArray, ndims_codomain::Val; kwargs...) +function one!!(A, ndims_codomain::Val; kwargs...) return one!!(FusionStyle(A), A, ndims_codomain; kwargs...) end -function one(style::FusionStyle, A::AbstractArray, ndims_codomain::Val; kwargs...) +function one(style::FusionStyle, A, ndims_codomain::Val; kwargs...) return one!!(style, copy(A), ndims_codomain; kwargs...) end -function one(A::AbstractArray, ndims_codomain::Val; kwargs...) +function one(A, ndims_codomain::Val; kwargs...) return one!!(copy(A), ndims_codomain; kwargs...) end diff --git a/src/linearbroadcasted.jl b/src/linearbroadcasted.jl index c848ae30..88440ec0 100644 --- a/src/linearbroadcasted.jl +++ b/src/linearbroadcasted.jl @@ -135,6 +135,9 @@ function Base.copy(a::Mul) end # copyto! for LinearBroadcasted dispatches to add!. +# Stays `AbstractArray`-bound: these overload `Base.copyto!`, so widening `dest` to `Any` +# collides with Base's own methods. A non-array destination (e.g. a wrapped `TensorMap`) +# gets an `AbstractTensorMap`-specific `copyto!` from the backend extension instead. function Base.copyto!(dest::AbstractArray, src::LinearBroadcasted) return add!(dest, src, true, false) end @@ -154,13 +157,13 @@ _compose_op(f, g) = f ∘ g # permutedimsopadd! for LinearBroadcasted subtypes. function permutedimsopadd!( - dest::AbstractArray, op, src::ScaledBroadcasted, perm, α::Number, β::Number + dest, op, src::ScaledBroadcasted, perm, α::Number, β::Number ) return permutedimsopadd!(dest, op, unscaled(src), perm, op(coeff(src)) * α, β) end function permutedimsopadd!( - dest::AbstractArray, op, src::AddBroadcasted, perm, α::Number, β::Number + dest, op, src::AddBroadcasted, perm, α::Number, β::Number ) args = addends(src) permutedimsopadd!(dest, op, first(args), perm, α, β) @@ -173,7 +176,7 @@ end # TODO: Replace with contractopadd! once that interface exists, # to avoid materializing the Mul intermediate. function permutedimsopadd!( - dest::AbstractArray, op, src::Mul, perm, α::Number, β::Number + dest, op, src::Mul, perm, α::Number, β::Number ) return permutedimsopadd!(dest, op, copy(src), perm, α, β) end @@ -199,8 +202,8 @@ linearbroadcasted(+, a, b) # AddBroadcasted(a, b) function linearbroadcasted end # Scaling: Number * AbstractArray -linearbroadcasted(::typeof(*), α::Number, a::AbstractArray) = ScaledBroadcasted(α, a) -linearbroadcasted(::typeof(*), a::AbstractArray, α::Number) = ScaledBroadcasted(α, a) +linearbroadcasted(::typeof(*), α::Number, a) = ScaledBroadcasted(α, a) +linearbroadcasted(::typeof(*), a, α::Number) = ScaledBroadcasted(α, a) # Scaling of ScaledBroadcasted: absorb coefficient. function linearbroadcasted(::typeof(*), α::Number, a::ScaledBroadcasted) return ScaledBroadcasted(α * coeff(a), unscaled(a)) @@ -208,7 +211,7 @@ end # Conjugation lowers to the `ConjArray` lazy wrapper. A scaled `ConjArray` (e.g. # `conj.(a) ./ β`) is handled by the generic `AbstractArray` scaling method above, since # `ConjArray <: AbstractArray`. -linearbroadcasted(::typeof(conj), a::AbstractArray) = conjed(a) +linearbroadcasted(::typeof(conj), a) = conjed(a) function linearbroadcasted(::typeof(conj), a::ScaledBroadcasted) return ScaledBroadcasted(conj(coeff(a)), linearbroadcasted(conj, unscaled(a))) end diff --git a/src/matricize.jl b/src/matricize.jl index b4eee31d..61d709ec 100644 --- a/src/matricize.jl +++ b/src/matricize.jl @@ -24,19 +24,13 @@ for a graded range). trivialrange(r::AbstractUnitRange) = trivialrange(typeof(r)) trivialrange(::Type{<:AbstractUnitRange}) = Base.OneTo(1) -# Default similar with bipartitioned axes: flatten to a plain tuple of axes. -# Downstream types (e.g., FusionTensor) can override to preserve bipartition. -function Base.similar(a::AbstractArray, T::Type, axes::BiTuple) - return similar(a, T, Tuple(axes)) -end - """ permutedimsop(op, src, perm_codomain, perm_domain) Non-mutating version of `bipermutedimsopadd!`: returns `op.(permutedims(src, (perm_codomain..., perm_domain...)))`. """ -function permutedimsop(op, src::AbstractArray, perm_codomain, perm_domain) +function permutedimsop(op, src, perm_codomain, perm_domain) dest = allocate_output(permutedimsop, op, src, perm_codomain, perm_domain) return bipermutedimsopadd!(dest, op, src, perm_codomain, perm_domain, true, false) end @@ -44,24 +38,26 @@ end # The output holds `op.(src)` permuted, so `op` applies to the axes too: `conj` dualizes a # graded axis (a no-op on a dense axis), `identity` leaves it unchanged, keeping axes and # data in sync. -function allocate_output(::typeof(permutedimsop), op, src::AbstractArray, perm_co, perm_do) +function allocate_output(::typeof(permutedimsop), op, src, perm_co, perm_do) T = Base.promote_op(op, eltype(src)) axes_co = map(i -> op(axes(src, i)), perm_co) axes_do = map(i -> op(axes(src, i)), perm_do) - return similar(src, T, BiTuple(axes_co, axes_do)) + # `axes_do` are in the stored/dualized convention (from `axes(src)`), so un-dualize them + # into `similar_map`'s codomain-facing convention. + return similar_map(src, T, axes_co, conj.(axes_do)) end -function bipermutedims(a::AbstractArray, perm1, perm2) +function bipermutedims(a, perm1, perm2) return permutedimsop(identity, a, perm1, perm2) end -function bipermutedims!(a_dest::AbstractArray, a_src::AbstractArray, perm1, perm2) +function bipermutedims!(a_dest, a_src, perm1, perm2) return bipermutedimsopadd!(a_dest, identity, a_src, perm1, perm2, true, false) end -function bipermutedims(a::AbstractArray, biperm::BiTuple) +function bipermutedims(a, biperm::BiTuple) return bipermutedims(a, biperm.t1, biperm.t2) end function bipermutedims!( - a_dest::AbstractArray, a_src::AbstractArray, biperm::BiTuple + a_dest, a_src, biperm::BiTuple ) return bipermutedims!(a_dest, a_src, biperm.t1, biperm.t2) end @@ -74,32 +70,32 @@ end # This is the primary function that should be overloaded for new fusion styles. # This assumes the permutation was already performed. function matricize( - style::FusionStyle, a::AbstractArray, ndims_codomain::Val + style::FusionStyle, a, ndims_codomain::Val ) return throw(MethodError(matricize, (style, a, ndims_codomain))) end -function matricize(a::AbstractArray, ndims_codomain::Val) +function matricize(a, ndims_codomain::Val) return matricize(FusionStyle(a), a, ndims_codomain) end -function matricize( - a::AbstractArray, +function matricizeperm( + a, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}} ) - return matricize(FusionStyle(a), a, perm_codomain, perm_domain) + return matricizeperm(FusionStyle(a), a, perm_codomain, perm_domain) end -# Thin wrapper around `matricizeop` with identity op — the actual matricization logic +# Thin wrapper around `matricizeopperm` with identity op — the actual matricization logic # (and the fusion-style overload point for folding ops into matricization) lives in -# `matricizeop`. -function matricize( - style::FusionStyle, a::AbstractArray, +# `matricizeopperm`. +function matricizeperm( + style::FusionStyle, a, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}} ) - return matricizeop(style, identity, a, perm_codomain, perm_domain) + return matricizeopperm(style, identity, a, perm_codomain, perm_domain) end # Process inputs such as `EllipsisNotation.Ellipsis`. -function to_permblocks(a::AbstractArray, permblocks::NTuple{2, Tuple{Vararg{Int}}}) +function to_permblocks(a, permblocks::NTuple{2, Tuple{Vararg{Int}}}) isperm((permblocks[1]..., permblocks[2]...)) || throw(ArgumentError("Invalid bipermutation")) return permblocks @@ -110,49 +106,49 @@ function tuplesetcomplement(t1::NTuple{N1}, t2::NTuple{N2}) where {N1, N2} return NTuple{N1 - N2}(setdiff(t1, t2)) end function to_permblocks( - a::AbstractArray, permblocks::Tuple{Tuple{Ellipsis}, Tuple{Vararg{Int}}} + a, permblocks::Tuple{Tuple{Ellipsis}, Tuple{Vararg{Int}}} ) permblocks1 = tuplesetcomplement(ntuple(identity, ndims(a)), permblocks[2]) return (permblocks1, permblocks[2]) end function to_permblocks( - a::AbstractArray, permblocks::Tuple{Tuple{Vararg{Int}}, Tuple{Ellipsis}} + a, permblocks::Tuple{Tuple{Vararg{Int}}, Tuple{Ellipsis}} ) permblocks2 = tuplesetcomplement(ntuple(identity, ndims(a)), permblocks[1]) return (permblocks[1], permblocks2) end -function matricize(a::AbstractArray, perm_codomain, perm_domain) - return matricize(FusionStyle(a), a, perm_codomain, perm_domain) +function matricizeperm(a, perm_codomain, perm_domain) + return matricizeperm(FusionStyle(a), a, perm_codomain, perm_domain) end -function matricize( - style::FusionStyle, a::AbstractArray, perm_codomain, perm_domain +function matricizeperm( + style::FusionStyle, a, perm_codomain, perm_domain ) - return matricize(style, a, to_permblocks(a, (perm_codomain, perm_domain))...) + return matricizeperm(style, a, to_permblocks(a, (perm_codomain, perm_domain))...) end -# ==================================== matricizeop ======================================= +# ================================== matricizeopperm ===================================== """ - matricizeop(op, a, perm_codomain, perm_domain) + matricizeopperm(op, a, perm_codomain, perm_domain) Matricize `a` with element-wise operation `op` folded in. Returns a matrix representing -`op.(matricize(a, perm_codomain, perm_domain))`. +`op.(matricizeperm(a, perm_codomain, perm_domain))`. Has "maybe alias" semantics: the result may be a view/wrapper aliasing `a` or a fresh copy, depending on the fusion style and array type. The caller should treat the result as read-only. """ -function matricizeop(op, a::AbstractArray, perm_codomain, perm_domain) - return matricizeop(FusionStyle(a), op, a, perm_codomain, perm_domain) +function matricizeopperm(op, a, perm_codomain, perm_domain) + return matricizeopperm(FusionStyle(a), op, a, perm_codomain, perm_domain) end -function matricizeop( - style::FusionStyle, op, a::AbstractArray, perm_codomain, perm_domain +function matricizeopperm( + style::FusionStyle, op, a, perm_codomain, perm_domain ) - return matricizeop(style, op, a, to_permblocks(a, (perm_codomain, perm_domain))...) + 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 `matricizeop` can skip the redundant permuted copy: +# 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 @@ -183,8 +179,8 @@ end # 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. -function matricizeop( - style::FusionStyle, op, a::AbstractArray, +function matricizeopperm( + style::FusionStyle, op, a, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}} ) ndims(a) == length(perm_codomain) + length(perm_domain) || @@ -201,54 +197,53 @@ function matricizeop( end # ==================================== unmatricize ======================================= -# This is the primary function that should be overloaded for new fusion styles. -function unmatricize( - style::FusionStyle, m::AbstractMatrix, - axes_codomain::Tuple{Vararg{AbstractUnitRange}}, - axes_domain::Tuple{Vararg{AbstractUnitRange}} - ) - return throw(MethodError(unmatricize, (style, m, axes_codomain, axes_domain))) +# Split form: `codomain_axes` and `domain_axes` are the destination axes for the codomain and +# domain groups. This is the primary overload point for new fusion styles. Permutation is +# handled separately by `unmatricizeperm`, so `unmatricize` never has to disambiguate axis +# tuples from permutation tuples regardless of how unconstrained `m` and the axes are. +function unmatricize(style::FusionStyle, m, codomain_axes, domain_axes) + return throw(MethodError(unmatricize, (style, m, codomain_axes, domain_axes))) end -function unmatricize( - m::AbstractMatrix, - axes_codomain::Tuple{Vararg{AbstractUnitRange}}, - axes_domain::Tuple{Vararg{AbstractUnitRange}} - ) - return unmatricize(FusionStyle(m), m, axes_codomain, axes_domain) +function unmatricize(m, codomain_axes, domain_axes) + return unmatricize(FusionStyle(m), m, codomain_axes, domain_axes) end -function unmatricize( - m::AbstractMatrix, axes_dest, - invperm1::Tuple{Vararg{Int}}, invperm2::Tuple{Vararg{Int}} +# Inverse-bipermutation form: split `axes_dest` into codomain/domain groups reordered by the +# inverse bipermutation, unmatricize in that order, then permute back. +function unmatricizeperm( + m, axes_dest, + invperm_codomain::Tuple{Vararg{Int}}, invperm_domain::Tuple{Vararg{Int}} ) - return unmatricize(FusionStyle(m), m, axes_dest, invperm1, invperm2) + return unmatricizeperm(FusionStyle(m), m, axes_dest, invperm_codomain, invperm_domain) end -function unmatricize( - style::FusionStyle, m::AbstractMatrix, axes_dest, +function unmatricizeperm( + style::FusionStyle, m, axes_dest, invperm_codomain::Tuple{Vararg{Int}}, invperm_domain::Tuple{Vararg{Int}} ) invbiperm = BiTuple(invperm_codomain, invperm_domain) length(axes_dest) == length(invbiperm) || throw(ArgumentError("axes do not match permutation")) - a12 = unmatricize(style, m, bipartition(axes_dest, invbiperm)...) + codomain_axes, domain_axes = bipartition(axes_dest, invbiperm) + a12 = unmatricize(style, m, codomain_axes, domain_axes) biperm_dest = BiTuple(Tuple(invperm(invbiperm)), Val(length_codomain(axes_dest))) return bipermutedims(a12, biperm_dest) end -function unmatricize!( - a_dest::AbstractArray, m::AbstractMatrix, +function unmatricizeperm!( + a_dest, m, invperm_codomain::Tuple{Vararg{Int}}, invperm_domain::Tuple{Vararg{Int}} ) - return unmatricize!(FusionStyle(m), a_dest, m, invperm_codomain, invperm_domain) + return unmatricizeperm!(FusionStyle(m), a_dest, m, invperm_codomain, invperm_domain) end -function unmatricize!( - style::FusionStyle, a_dest::AbstractArray, m::AbstractMatrix, +function unmatricizeperm!( + style::FusionStyle, a_dest, m, invperm_codomain::Tuple{Vararg{Int}}, invperm_domain::Tuple{Vararg{Int}} ) invbiperm = BiTuple(invperm_codomain, invperm_domain) ndims(a_dest) == length(invbiperm) || throw(ArgumentError("destination does not match permutation")) - a_perm = unmatricize(style, m, bipartition(axes(a_dest), invbiperm)...) + codomain_axes, domain_axes = bipartition(axes(a_dest), invbiperm) + a_perm = unmatricize(style, m, codomain_axes, domain_axes) biperm_dest = BiTuple(Tuple(invperm(invbiperm)), Val(length_codomain(axes(a_dest)))) return bipermutedims!(a_dest, a_perm, biperm_dest) end @@ -256,7 +251,7 @@ end # Defaults to ReshapeFusion, a simple reshape struct ReshapeFusion <: FusionStyle end FusionStyle(::Type{<:AbstractArray}) = ReshapeFusion() -function matricize(::ReshapeFusion, a::AbstractArray, ndims_codomain::Val) +function matricize(::ReshapeFusion, 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) @@ -272,10 +267,7 @@ function matricizekind( isidentityperm((perm_domain..., perm_codomain...)) && return TransposeMatricizeKind return PermuteMatricizeKind end -function unmatricize( - style::ReshapeFusion, m::AbstractMatrix, - axes_codomain::Tuple{Vararg{AbstractUnitRange}}, - axes_domain::Tuple{Vararg{AbstractUnitRange}} - ) - return reshape(m, (axes_codomain..., axes_domain...)) +# A dense reshape ignores the codomain/domain split: it just reshapes to the concatenated axes. +function unmatricize(style::ReshapeFusion, m, codomain_axes, domain_axes) + return reshape(m, (codomain_axes..., domain_axes...)) end diff --git a/src/matrixfunctions.jl b/src/matrixfunctions.jl index 7628cc58..e0676aa2 100644 --- a/src/matrixfunctions.jl +++ b/src/matrixfunctions.jl @@ -33,17 +33,18 @@ const MATRIX_FUNCTIONS = [ for f in MATRIX_FUNCTIONS @eval begin - function $f(style::FusionStyle, a::AbstractArray, ndims_codomain::Val; kwargs...) + function $f(style::FusionStyle, a, ndims_codomain::Val; kwargs...) a_mat = matricize(style, a, ndims_codomain) fa_mat = Base.$f(a_mat; kwargs...) - return unmatricize(style, fa_mat, bipartition(axes(a), ndims_codomain)...) + codomain_axes, domain_axes = bipartition(axes(a), ndims_codomain) + return unmatricize(style, fa_mat, codomain_axes, domain_axes) end - function $f(a::AbstractArray, ndims_codomain::Val; kwargs...) + function $f(a, ndims_codomain::Val; kwargs...) return $f(FusionStyle(a), a, ndims_codomain; kwargs...) end function $f( - style::FusionStyle, a::AbstractArray, + style::FusionStyle, a, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs... ) @@ -51,7 +52,7 @@ for f in MATRIX_FUNCTIONS return $f(style, a_perm, Val(length(perm_codomain)); kwargs...) end function $f( - a::AbstractArray, + a, perm_codomain::Tuple{Vararg{Int}}, perm_domain::Tuple{Vararg{Int}}; kwargs... ) @@ -60,7 +61,7 @@ for f in MATRIX_FUNCTIONS end function $f( - style::FusionStyle, a::AbstractArray, + style::FusionStyle, a, labels_a, labels_codomain, labels_domain; kwargs... ) perm_codomain, perm_domain = @@ -68,7 +69,7 @@ for f in MATRIX_FUNCTIONS return $f(style, a, perm_codomain, perm_domain; kwargs...) end function $f( - a::AbstractArray, + a, labels_a, labels_codomain, labels_domain; kwargs... ) perm_codomain, perm_domain = diff --git a/src/permutedimsadd.jl b/src/permutedimsadd.jl index d791df6e..82d1780b 100644 --- a/src/permutedimsadd.jl +++ b/src/permutedimsadd.jl @@ -2,12 +2,30 @@ using Strided: Strided using StridedViews: StridedViews as SV """ - permuteddims(a::AbstractArray, perm) + permuteddims(a, perm) -Lazy `permutedims`, defaulting to a `Base.PermutedDimsArray` view. This is an extension -hook: downstream array types can overload it to return a custom lazy permuted-dims type. +Lazy `permutedims`. For an `AbstractArray` this is a `Base.PermutedDimsArray` view; for any +other operand it is a generic [`PermutedDims`](@ref) node. This is an extension hook: +downstream types can overload it to return their own lazy permuted-dims type. """ permuteddims(a::AbstractArray, perm) = PermutedDimsArray(a, perm) +permuteddims(a, perm) = PermutedDims(a, perm) + +""" + PermutedDims(parent, perm) + +Generic lazy permuted-dims wrapper for operands that are not `AbstractArray`s (so +`Base.PermutedDimsArray` does not apply), modeled on `PermutedDimsArray`: it records `parent` +and the permutation `perm` without materializing anything. TensorAlgebra never indexes into +it. It exists so that a downstream extension's `permuteddims` can hand a lazily permuted +non-array operand to `bipermutedimsopadd!`, whose absorption method composes `perm` into the +outer bipermutation and forwards to a single leaf call on `parent`. +""" +struct PermutedDims{P, Perm} + parent::P + perm::Perm +end +Base.parent(a::PermutedDims) = a.parent # ---------------------------------------------------------------------------- # # bipermutedimsopadd! — the primary materialization primitive @@ -20,7 +38,7 @@ function bipermutedimsopadd! end # of which act on axes: `conj` dualizes a graded axis (and is a no-op on a dense axis), # `identity` leaves it unchanged. function check_input( - ::typeof(bipermutedimsopadd!), dest::AbstractArray, op, src::AbstractArray, + ::typeof(bipermutedimsopadd!), dest, op, src, perm_codomain, perm_domain ) op === identity || op === conj || @@ -51,7 +69,7 @@ element-wise, permutes, then accumulates via broadcasting with Strided.jl optimization when possible. """ function bipermutedimsopadd!( - dest::AbstractArray, op, src::AbstractArray, + dest, op, src, perm_codomain, perm_domain, α::Number, β::Number ) @@ -82,9 +100,13 @@ function _opadd!(dest::AbstractArray, op, src::AbstractArray, α, β) end _permuteddims_perm(::PermutedDimsArray{<:Any, <:Any, perm}) where {perm} = perm +_permuteddims_perm(a::PermutedDims) = a.perm +# Both the dense `PermutedDimsArray` view and the generic `PermutedDims` node absorb the same +# way: compose the wrapper's permutation `w` into the outer bipermutation and forward to a +# single leaf `bipermutedimsopadd!` on the parent. function bipermutedimsopadd!( - dest::AbstractArray, op, src::PermutedDimsArray, + dest, op, src::Union{PermutedDimsArray, PermutedDims}, perm_codomain, perm_domain, α::Number, β::Number ) @@ -111,7 +133,7 @@ bipartitioned permutation version; this flat-permutation overload forwards to it with `perm_domain = ()`. """ function permutedimsopadd!( - dest::AbstractArray, op, src::AbstractArray, perm, α::Number, β::Number + dest, op, src, perm, α::Number, β::Number ) return bipermutedimsopadd!(dest, op, src, perm, (), α, β) end @@ -126,7 +148,7 @@ end `dest = β * dest + α * permutedims(src, perm)`. """ function permutedimsadd!( - dest::AbstractArray, src::AbstractArray, perm, α::Number, β::Number + dest, src, perm, α::Number, β::Number ) return permutedimsopadd!(dest, identity, src, perm, α, β) end @@ -136,7 +158,7 @@ end `dest = β * dest + α * src`. """ -function add!(dest::AbstractArray, src, α::Number, β::Number) +function add!(dest, src, α::Number, β::Number) return permutedimsopadd!(dest, identity, src, ntuple(identity, ndims(src)), α, β) end @@ -145,4 +167,4 @@ end `dest .+= src`. """ -add!(dest::AbstractArray, src) = add!(dest, src, true, true) +add!(dest, src) = add!(dest, src, true, true) diff --git a/src/similar_map.jl b/src/similar_map.jl index da408b76..d527564c 100644 --- a/src/similar_map.jl +++ b/src/similar_map.jl @@ -1,10 +1,13 @@ """ similar_map(prototype, [T,] codomain_axes, domain_axes) -> M -Allocate an array shaped as a linear map from `domain_axes` to -`codomain_axes` with element type `T` (defaulting to `eltype(prototype)`), -using `prototype` to determine the array backend. Defaults to -`similar(prototype, T, (codomain_axes..., conj.(domain_axes)...))`. +Allocate an array shaped as a linear map from `domain_axes` to `codomain_axes` +with element type `T` (defaulting to `eltype(prototype)`), using `prototype` to +determine the array backend. The domain axes are given un-dualized (codomain +facing) and stored dual, so the default is +`similar(prototype, T, (codomain_axes..., conj.(domain_axes)...))`. `conj` +dualizes a graded axis and is a no-op on a dense axis. Backends with map-shaped +storage (e.g. a `TensorMap`) overload this to build the codomain/domain directly. # Examples diff --git a/test/Project.toml b/test/Project.toml index 527c0bc8..a8e405f2 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -36,7 +36,7 @@ Random = "1.10" SafeTestsets = "0.1" StableRNGs = "1.0.2" Suppressor = "0.2" -TensorAlgebra = "0.14" +TensorAlgebra = "0.15" TensorOperations = "5.1.4" Test = "1.10" TestExtras = "0.3.1" diff --git a/test/test_basics.jl b/test/test_basics.jl index 2a5426a4..100e65e6 100644 --- a/test/test_basics.jl +++ b/test/test_basics.jl @@ -2,8 +2,8 @@ import TensorAlgebra using EllipsisNotation: var".." using StableRNGs: StableRNG using TensorAlgebra: BiTuple, ContractAlgorithm, bipermutedims, bipermutedims!, contract, - contract!, contractadd!, length_codomain, length_domain, matricize, unmatricize, - unmatricize! + contract!, contractadd!, length_codomain, length_domain, matricizeperm, unmatricize, + unmatricizeperm, unmatricizeperm! using TensorOperations: TensorOperations using Test: @test, @test_broken, @test_throws, @testset @@ -54,70 +54,70 @@ TensorAlgebra.label_type(::Type{OptInLabel}) = Int @testset "matricize (eltype=$elt)" for elt in elts a = randn(elt, 2, 3, 4, 5) - a_fused = matricize(a, (1, 2), (3, 4)) + a_fused = matricizeperm(a, (1, 2), (3, 4)) @test eltype(a_fused) === elt @test a_fused ≈ reshape(a, 6, 20) - a_fused = matricize(a, (3, 1), (2, 4)) + a_fused = matricizeperm(a, (3, 1), (2, 4)) @test eltype(a_fused) === elt @test a_fused ≈ reshape(permutedims(a, (3, 1, 2, 4)), (8, 15)) - a_fused = matricize(a, (3, 1, 2), (4,)) + a_fused = matricizeperm(a, (3, 1, 2), (4,)) @test eltype(a_fused) === elt @test a_fused ≈ reshape(permutedims(a, (3, 1, 2, 4)), (24, 5)) - a_fused = matricize(a, (..,), (3, 1)) + a_fused = matricizeperm(a, (..,), (3, 1)) @test eltype(a_fused) === elt @test a_fused ≈ reshape(permutedims(a, (2, 4, 3, 1)), (15, 8)) - a_fused = matricize(a, (3, 1), (..,)) + a_fused = matricizeperm(a, (3, 1), (..,)) @test eltype(a_fused) === elt @test a_fused ≈ reshape(permutedims(a, (3, 1, 2, 4)), (8, 15)) - a_fused = matricize(a, (), (..,)) + a_fused = matricizeperm(a, (), (..,)) @test eltype(a_fused) === elt @test a_fused ≈ reshape(a, (1, 120)) - a_fused = matricize(a, (..,), ()) + a_fused = matricizeperm(a, (..,), ()) @test eltype(a_fused) === elt @test a_fused ≈ reshape(a, (120, 1)) - @test_throws MethodError matricize(a, (1, 2), (3,), (4,)) - @test_throws MethodError matricize(a, (1, 2, 3, 4)) - @test_throws ArgumentError matricize(a, (1, 2), (3,)) + @test_throws MethodError matricizeperm(a, (1, 2), (3,), (4,)) + @test_throws MethodError matricizeperm(a, (1, 2, 3, 4)) + @test_throws ArgumentError matricizeperm(a, (1, 2), (3,)) v = ones(elt, 2) - a_fused = matricize(v, (1,), ()) + a_fused = matricizeperm(v, (1,), ()) @test eltype(a_fused) === elt @test a_fused ≈ ones(elt, 2, 1) - a_fused = matricize(v, (), (1,)) + a_fused = matricizeperm(v, (), (1,)) @test eltype(a_fused) === elt @test a_fused ≈ ones(elt, 1, 2) - a_fused = matricize(ones(elt), (), ()) + a_fused = matricizeperm(ones(elt), (), ()) @test eltype(a_fused) === elt @test a_fused ≈ ones(elt, 1, 1) end - @testset "matricizeop (eltype=$elt)" for elt in elts + @testset "matricizeopperm (eltype=$elt)" for elt in elts rng = StableRNG(123) a = randn(rng, elt, 2, 3, 4) # identity op: should match matricize exactly - m = TensorAlgebra.matricizeop(identity, a, (1,), (2, 3)) - m_ref = matricize(a, (1,), (2, 3)) + m = TensorAlgebra.matricizeopperm(identity, a, (1,), (2, 3)) + m_ref = matricizeperm(a, (1,), (2, 3)) @test m ≈ m_ref - m = TensorAlgebra.matricizeop(identity, a, (3, 1), (2,)) - m_ref = matricize(a, (3, 1), (2,)) + m = TensorAlgebra.matricizeopperm(identity, a, (3, 1), (2,)) + m_ref = matricizeperm(a, (3, 1), (2,)) @test m ≈ m_ref - m = TensorAlgebra.matricizeop(identity, a, (2, 3), (1,)) - m_ref = matricize(a, (2, 3), (1,)) + m = TensorAlgebra.matricizeopperm(identity, a, (2, 3), (1,)) + m_ref = matricizeperm(a, (2, 3), (1,)) @test m ≈ m_ref # conj op - m = TensorAlgebra.matricizeop(conj, a, (1,), (2, 3)) - m_ref = conj.(matricize(a, (1,), (2, 3))) + m = TensorAlgebra.matricizeopperm(conj, a, (1,), (2, 3)) + m_ref = conj.(matricizeperm(a, (1,), (2, 3))) @test m ≈ m_ref - m = TensorAlgebra.matricizeop(conj, a, (3, 1), (2,)) - m_ref = conj.(matricize(a, (3, 1), (2,))) + m = TensorAlgebra.matricizeopperm(conj, a, (3, 1), (2,)) + m_ref = conj.(matricizeperm(a, (3, 1), (2,))) @test m ≈ m_ref end @@ -130,7 +130,7 @@ TensorAlgebra.label_type(::Type{OptInLabel}) = Int @test eltype(a) === elt @test a ≈ a0 - a = unmatricize(m, axes0, (1, 2), (3, 4)) + a = unmatricizeperm(m, axes0, (1, 2), (3, 4)) @test eltype(a) === elt @test a ≈ a0 @@ -139,21 +139,21 @@ TensorAlgebra.label_type(::Type{OptInLabel}) = Int invperm_codomain = (3, 2) invperm_domain = (4, 1) perm = (4, 2, 1, 3) - a = unmatricize(m, map(i -> axes0[i], perm), invperm_codomain, invperm_domain) + a = unmatricizeperm(m, map(i -> axes0[i], perm), invperm_codomain, invperm_domain) @test eltype(a) === elt @test a ≈ permutedims(a0, perm) a = similar(a0) - unmatricize!(a, m, (1, 2), (3, 4)) + unmatricizeperm!(a, m, (1, 2), (3, 4)) @test a ≈ a0 - m1 = matricize(a0, perm_codomain, perm_domain) - a = unmatricize(m1, axes0, perm_codomain, perm_domain) + m1 = matricizeperm(a0, perm_codomain, perm_domain) + a = unmatricizeperm(m1, axes0, perm_codomain, perm_domain) @test a ≈ a0 a1 = permutedims(a0, perm) a = similar(a1) - unmatricize!(a, m, invperm_codomain, invperm_domain) + unmatricizeperm!(a, m, invperm_codomain, invperm_domain) @test a ≈ a1 a = unmatricize(m, (), axes0) @@ -169,8 +169,8 @@ TensorAlgebra.label_type(::Type{OptInLabel}) = Int @test a isa Array{elt, 0} @test a[] == m[1, 1] - @test_throws ArgumentError unmatricize(m, (), (1, 2), (3,)) - @test_throws ArgumentError unmatricize!(m, m, (1, 2), (3,)) + @test_throws ArgumentError unmatricizeperm(m, (), (1, 2), (3,)) + @test_throws ArgumentError unmatricizeperm!(m, m, (1, 2), (3,)) end alg_tensoroperations = ContractAlgorithm(TensorOperations.StridedBLAS()) diff --git a/test/test_exports.jl b/test/test_exports.jl index 96b725a7..1ccdd1c0 100644 --- a/test/test_exports.jl +++ b/test/test_exports.jl @@ -34,7 +34,7 @@ using Test: @test, @testset append!( exports, [ - :biperm, :bipartition, :contractopadd!, :label_type, :matricizeop, + :biperm, :bipartition, :contractopadd!, :label_type, :matricizeopperm, :to_range, :zero!, :scale!, :permuteddims, :conjed, :ConjArray, ] ) diff --git a/test/test_matricize.jl b/test/test_matricize.jl index fdd8f0f4..7294ef09 100644 --- a/test/test_matricize.jl +++ b/test/test_matricize.jl @@ -1,7 +1,7 @@ using LinearAlgebra: Transpose using StableRNGs: StableRNG using TensorAlgebra: TensorAlgebra, PermuteMatricizeKind, ReshapeFusion, - ReshapeMatricizeKind, TransposeMatricizeKind, matricize, matricizeop + ReshapeMatricizeKind, TransposeMatricizeKind, matricizeopperm, matricizeperm using Test: @test, @testset # A non-`ReshapeFusion` style, to check the always-safe generic fallback. @@ -37,30 +37,30 @@ end @test TensorAlgebra.matricizekind(DummyFusion(), (3, 1), (2,)) == PermuteMatricizeKind end -@testset "maybe-view matricizeop (eltype=$elt)" for elt in (Float64, ComplexF64) +@testset "maybe-view matricizeopperm (eltype=$elt)" for elt in (Float64, ComplexF64) a = randn(StableRNG(123), elt, 2, 3, 4) # Reshape branch: correct values and a view aliasing `a`. - m = matricize(a, (1,), (2, 3)) + 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 = matricize(a, (2, 3), (1,)) + 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 = matricize(a, (3, 1), (2,)) + m = matricizeperm(a, (3, 1), (2,)) @test m ≈ matricize_ref(a, (3, 1), (2,)) @test !Base.mightalias(m, a) # `conj` cannot ride a view, so it copies even on the reshape/transpose patterns. - m = matricizeop(conj, a, (1,), (2, 3)) + m = matricizeopperm(conj, a, (1,), (2, 3)) @test m ≈ conj.(matricize_ref(a, (1,), (2, 3))) @test !Base.mightalias(m, a) - m = matricizeop(conj, a, (2, 3), (1,)) + m = matricizeopperm(conj, a, (2, 3), (1,)) @test m ≈ conj.(matricize_ref(a, (2, 3), (1,))) @test !Base.mightalias(m, a) end @@ -70,19 +70,19 @@ end # Reshape view tracks an in-place update of `a`. a = randn(rng, 2, 3, 4) - m = matricize(a, (1,), (2, 3)) + 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 = matricize(a, (2, 3), (1,)) + 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 = matricize(a, (3, 1), (2,)) + m = matricizeperm(a, (3, 1), (2,)) snapshot = copy(m) a .= a .+ 1 @test m == snapshot diff --git a/test/test_permutedimsadd.jl b/test/test_permutedimsadd.jl index 9e417bfa..58383690 100644 --- a/test/test_permutedimsadd.jl +++ b/test/test_permutedimsadd.jl @@ -1,9 +1,14 @@ using Adapt: adapt using JLArrays: JLArray -using TensorAlgebra: - ConjArray, add!, bipermutedimsopadd!, conjed, permutedimsadd!, permutedimsopadd! +using TensorAlgebra: ConjArray, PermutedDims, add!, bipermutedimsopadd!, conjed, + permuteddims, permutedimsadd!, permutedimsopadd! using Test: @test, @testset +# A non-`AbstractArray` operand, to check that `permuteddims` falls back to `PermutedDims`. +struct NotAnArray{P} + parent::P +end + @testset "[permutedims]add!" begin @testset "add!(b, a, α, β) (arraytype=$arrayt)" for arrayt in (Array, JLArray) dev = adapt(arrayt) @@ -68,6 +73,37 @@ using Test: @test, @testset end end end + @testset "permuteddims dispatch" begin + a = randn(2, 3, 4) + # An `AbstractArray` gets a `Base.PermutedDimsArray` view. + @test permuteddims(a, (3, 1, 2)) isa PermutedDimsArray + @test permuteddims(a, (3, 1, 2)) == permutedims(a, (3, 1, 2)) + # Any other operand gets a generic `PermutedDims` node wrapping it unchanged. + x = NotAnArray(a) + p = permuteddims(x, (3, 1, 2)) + @test p isa PermutedDims + @test parent(p) === x + end + @testset "bipermutedimsopadd! unwraps PermutedDims src (arraytype=$arrayt)" for arrayt in + ( + Array, + JLArray, + ) + dev = adapt(arrayt) + parent = dev(randn(2, 3, 4, 5)) + w = (3, 1, 4, 2) + src = PermutedDims(parent, w) + for (pc, pd) in (((1, 2, 3, 4), ()), ((2, 4), (1, 3)), ((3, 1), (2, 4))) + perm = (pc..., pd...) + ref = permutedims(permutedims(parent, w), perm) + for β in (0, 3) + dest = dev(randn(size(ref)...)) + dest′ = copy(dest) + bipermutedimsopadd!(dest′, identity, src, pc, pd, 2, β) + @test dest′ ≈ β * dest + 2 * ref + end + end + end @testset "bipermutedimsopadd! unwraps ConjArray src (arraytype=$arrayt)" for arrayt in ( Array,