Rework the matricize, contract, and factorization surfaces

Project.toml

@@ -1,13 +1,12 @@
 name = "TensorAlgebra"
 uuid = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
-version = "0.20.1"
+version = "0.21.0"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
 
 [workspace]
 projects = ["benchmark", "dev", "docs", "examples", "test"]
 
 [deps]
-EllipsisNotation = "da5c29d0-fa7d-589e-88eb-ea29b0a81949"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MatrixAlgebraKit = "6c742aac-3347-4629-af66-fc926824e5e4"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -29,7 +28,6 @@ TensorAlgebraTensorKitSectorsExt = "TensorKitSectors"
 TensorAlgebraTensorOperationsExt = "TensorOperations"
 
 [compat]
-EllipsisNotation = "1.8"
 LinearAlgebra = "1.10"
 MatrixAlgebraKit = "0.6"
 Mooncake = "0.4.202, 0.5"

docs/Project.toml

@@ -11,4 +11,4 @@ path = ".."
 Documenter = "1.8.1"
 ITensorFormatter = "0.2.27"
 Literate = "2.20.1"
-TensorAlgebra = "0.20"
+TensorAlgebra = "0.21"

examples/Project.toml

@@ -5,4 +5,4 @@ TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
 path = ".."
 
 [compat]
-TensorAlgebra = "0.20"
+TensorAlgebra = "0.21"

ext/TensorAlgebraMooncakeExt/TensorAlgebraMooncakeExt.jl

@@ -3,7 +3,7 @@ module TensorAlgebraMooncakeExt
 using Mooncake: Mooncake, @zero_derivative, DefaultCtx
 using TensorAlgebra: BiTuple, ContractAlgorithm, allocate_output, biperm, biperms,
     check_input, contract, contract!, contract_labels, decode_contraction_labels,
-    default_contract_algorithm, encode_contraction_labels, select_contract_algorithm
+    default_algorithm, encode_contraction_labels, select_algorithm
 
 Mooncake.tangent_type(::Type{<:BiTuple}) = Mooncake.NoTangent
 Mooncake.tangent_type(::Type{<:ContractAlgorithm}) = Mooncake.NoTangent
@@ -32,7 +32,7 @@ Mooncake.tangent_type(::Type{<:ContractAlgorithm}) = Mooncake.NoTangent
 @zero_derivative DefaultCtx Tuple{typeof(contract_labels), Any, Any, Any, Any}
 @zero_derivative DefaultCtx Tuple{typeof(encode_contraction_labels), Any, Any}
 @zero_derivative DefaultCtx Tuple{typeof(decode_contraction_labels), Any, Any, Any}
-@zero_derivative DefaultCtx Tuple{typeof(default_contract_algorithm), Any, Any}
-@zero_derivative DefaultCtx Tuple{typeof(select_contract_algorithm), Any, Any, Any}
+@zero_derivative DefaultCtx Tuple{typeof(default_algorithm), Any, Any}
+@zero_derivative DefaultCtx Tuple{typeof(select_algorithm), Any, Any, Any}
 
 end

ext/TensorAlgebraTensorKitExt.jl

@@ -243,44 +243,47 @@ end
 struct TensorKitMatricize <: TensorAlgebra.MatricizeStyle end
 TensorAlgebra.MatricizeStyle(::Type{<:AbstractTensorMap}) = TensorKitMatricize()
 
-# `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
+# `permute` at the tensor's own codomain/domain split is trivial and returns `t` itself, so that
+# split is the one memory-sharing matricization (TensorKit's own `has_shared_permute` notion). Any
+# other split, or a folded `conj`, regroups into a fresh `TensorMap`.
+function TensorAlgebra.is_output_view(
+        ::typeof(TensorAlgebra.matricizeop), ::TensorKitMatricize, op,
+        t::AbstractTensorMap, perm_codomain, perm_domain
+    )
+    return op === identity &&
+        TensorAlgebra.isidentitybiperm(perm_codomain, perm_domain) &&
+        length(perm_codomain) == numout(t)
 end
-TensorAlgebra.ismatricizeview(::TensorKitMatricize, ::AbstractTensorMap, ::Val) = false
-function TensorAlgebra.matricizeview(
-        ::TensorKitMatricize, t::AbstractTensorMap{<:Any, <:Any, K}, ::Val{K}
-    ) where {K}
+function TensorAlgebra.matricizeopview(
+        ::TensorKitMatricize, op, t::AbstractTensorMap, perm_codomain, perm_domain
+    )
     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)));
-        copy = true
+# A `TensorMap`'s matricization is a regrouping of its indices, so the destination is the same
+# `TensorMap` the plain permuted-add would allocate. That one already dualizes each space under
+# `op === conj`, which is what `bipermutedimsopadd!` needs: `tensoradd!` realizes a conjugation
+# by adjointing its source, so a destination over the undualized space does not match it.
+function TensorAlgebra.allocate_output(
+        ::typeof(TensorAlgebra.matricizeop), ::TensorKitMatricize, op,
+        t::AbstractTensorMap, perm_codomain, perm_domain
+    )
+    return TensorAlgebra.allocate_output(
+        TensorAlgebra.permutedimsop, op, t, perm_codomain, perm_domain
     )
 end
-
-# The identity fill on the regrouped map is TensorKit's own `one!` (MatrixAlgebraKit's
-# `one!` speaks `AbstractMatrix` only).
-function TensorAlgebra.one!!(
-        style::TensorKitMatricize, A::AbstractTensorMap, ndims_codomain::Val; kwargs...
+function TensorAlgebra.matricizeop!(
+        dest::AbstractTensorMap, ::TensorKitMatricize, op,
+        t::AbstractTensorMap, perm_codomain, perm_domain
+    )
+    return TensorAlgebra.bipermutedimsopadd!(
+        dest, op, t, perm_codomain, perm_domain, true, false
     )
-    return TensorKit.one!(TensorAlgebra.matricize(style, A, ndims_codomain))
 end
 
-# `unmatricize` reconstructs the codomain/domain axes from the matrix `m`. A `TensorMap` already
-# is the linear map its space describes, so the only valid request is the one whose codomain/domain
-# split matches `m`'s own space, and `unmatricize` returns `m` unchanged. The domain axes arrive
-# codomain-facing (un-dualized), which is exactly TensorKit's domain convention, so they build the
-# domain `ProductSpace` directly.
+# TensorKit owns its own `one!` generic and forwards only the `AbstractMatrix` case to
+# MatrixAlgebraKit, so the identity fill on a regrouped `TensorMap` goes through TensorKit.
+TensorAlgebra.MatrixAlgebra.one!(t::AbstractTensorMap) = TensorKit.one!(t)
+
 function TensorAlgebra.unmatricize(
         ::TensorKitMatricize, m::AbstractTensorMap, axes_codomain, axes_domain
     )
@@ -296,7 +299,7 @@ end
 # already implements through its TensorOperations interface. Route the generic `contract`
 # there: `zero!` clears the `similar_map`-allocated destination, and the default algorithm
 # hands the in-place contraction to the TensorOperations backend (see the TensorOperations
-# extension's `contractopadd!`).
+# extension's `contractpermopadd!`).
 TensorAlgebra.zero!(t::AbstractTensorMap) = VectorInterface.zerovector!(t)
 
 # A `TensorMap` is not an `AbstractArray`, so the generic in-place `TensorAlgebra` operations
@@ -310,8 +313,9 @@ function TensorAlgebra.add!(
     return VectorInterface.add!(y, x, α, β)
 end
 
-function TensorAlgebra.default_contract_algorithm(
-        ::Type{<:AbstractTensorMap}, ::Type{<:AbstractTensorMap}
+function TensorAlgebra.default_algorithm(
+        ::typeof(TensorAlgebra.contract!),
+        ::Type{<:Tuple{AbstractTensorMap, AbstractTensorMap, AbstractTensorMap}}
     )
     return TensorAlgebra.ContractAlgorithm(TO.DefaultBackend())
 end

ext/TensorAlgebraTensorOperationsExt/TensorAlgebraTensorOperationsExt.jl

@@ -1,63 +1,30 @@
 module TensorAlgebraTensorOperationsExt
 
-using TensorAlgebra: TensorAlgebra as TA, TensorOperationsAlgorithm
+using TensorAlgebra: TensorAlgebra as TA, TensorOperationsContract
 using TensorOperations: TensorOperations as TO
 
-# `TensorOperationsAlgorithm` stores `nothing` to mean "TensorOperations' default"; resolve
+# `TensorOperationsContract` stores `nothing` to mean "TensorOperations' default"; resolve
 # those here, where the defaults can be named.
-function backend(algorithm::TensorOperationsAlgorithm)
+function backend(algorithm::TensorOperationsContract)
     return @something algorithm.backend TO.DefaultBackend()
 end
-function allocator(algorithm::TensorOperationsAlgorithm)
+function allocator(algorithm::TensorOperationsContract)
     return @something algorithm.allocator TO.DefaultAllocator()
 end
 
 # Construct via the `ContractAlgorithm` public constructor seam as well.
-TA.ContractAlgorithm(backend::TO.AbstractBackend) = TensorOperationsAlgorithm(; backend)
+function TA.ContractAlgorithm(backend::TO.AbstractBackend)
+    return TensorOperationsContract(; backend)
+end
 function TA.ContractAlgorithm(backend::TO.AbstractBackend, allocator)
-    return TensorOperationsAlgorithm(; backend, allocator)
+    return TensorOperationsContract(; backend, allocator)
 end
 
 # Using TensorOperations backends as TensorAlgebra implementations
 # ----------------------------------------------------------------
 
-# not in-place
-function TA.contract(
-        algorithm::TensorOperationsAlgorithm,
-        perm_dest_codomain, perm_dest_domain,
-        a1::AbstractArray, perm1_codomain, perm1_domain,
-        a2::AbstractArray, perm2_codomain, perm2_domain
-    )
-    permblocks1 = Tuple.((perm1_codomain, perm1_domain))
-    permblocks2 = Tuple.((perm2_codomain, perm2_domain))
-    permblocks_dest = Tuple.((perm_dest_codomain, perm_dest_domain))
-    conj1, conj2 = false, false
-    α = true
-    return TO.tensorcontract(
-        a1, permblocks1, conj1, a2, permblocks2, conj2,
-        permblocks_dest, α, backend(algorithm), allocator(algorithm)
-    )
-end
-
-function TA.contract(
-        algorithm::TensorOperationsAlgorithm,
-        labels_dest,
-        a1::AbstractArray, labels1,
-        a2::AbstractArray, labels2
-    )
-    permblocks1, permblocks2, permblocks_dest =
-        TO.contract_indices(labels1, labels2, labels_dest)
-    conj1, conj2 = false, false
-    α = true
-    return TO.tensorcontract(
-        a1, permblocks1, conj1, a2, permblocks2, conj2,
-        permblocks_dest, α, backend(algorithm), allocator(algorithm)
-    )
-end
-
-# in-place
-function TA.contractopadd!(
-        algorithm::TensorOperationsAlgorithm,
+function TA.contractpermopadd!(
+        algorithm::TensorOperationsContract,
         a_dest, perm_dest_codomain, perm_dest_domain,
         op1, a1, perm1_codomain, perm1_domain,
         op2, a2, perm2_codomain, perm2_domain,
@@ -89,7 +56,7 @@ function TO.tensorcontract!(
     )
     op1 = conj1 ? conj : identity
     op2 = conj2 ? conj : identity
-    return TA.contractopadd!(
+    return TA.contractpermopadd!(
         backend,
         a_dest, permblocks_dest...,
         op1, a1, permblocks1...,

src/MatrixAlgebra.jl

@@ -1,19 +1,28 @@
 module MatrixAlgebra
 
-export gram_eigh_full,
-    gram_eigh_full_with_pinv,
-    invsqrt_diag_safe,
+export invsqrt_diag_safe,
     invsqrth_safe,
+    one!,
     pow_diag_safe,
     pow_diag_safe!,
     powh_safe,
     sqrt_diag_safe,
-    sqrth_safe,
-    sqrth_invsqrth_safe
+    sqrth_invsqrth_safe,
+    sqrth_safe
 
 using LinearAlgebra: LinearAlgebra, Diagonal, isdiag, norm
 using MatrixAlgebraKit: MatrixAlgebraKit as MAK
 
+"""
+    MatrixAlgebra.one!(m) -> m
+
+Fill `m` with the identity in place. The matrix-level identity fill the tensor-level
+`TensorAlgebra.one`/`one!` bottom out on, and the customization point a backend overloads when its
+matricization is a type `MatrixAlgebraKit.one!` does not handle (TensorKit owns its own `one!`
+generic rather than extending MatrixAlgebraKit's, so a `TensorMap` fills through that).
+"""
+one!(m) = MAK.one!(m)
+
 function _clamp_kwargs_doc(arg::AbstractString)
     return join(
         (
@@ -193,92 +202,6 @@ function sqrth_invsqrth_safe(M; alg = nothing, kwargs...)
         V * pow_diag_safe(D, -1 // 2; kwargs...) * V'
 end
 
-for (gram, gram_with_pinv, eigh_full) in (
-        (:gram_eigh_full, :gram_eigh_full_with_pinv, :eigh_full),
-        (:gram_eigh_full!!, :gram_eigh_full_with_pinv!!, :eigh_full!),
-    )
-    @eval begin
-        function $gram(A::AbstractMatrix; alg = nothing, kwargs...)
-            D, V = MAK.$eigh_full(A; alg)
-            return V * sqrth_safe(D; kwargs...)
-        end
-        function $gram_with_pinv(A::AbstractMatrix; alg = nothing, kwargs...)
-            D, V = MAK.$eigh_full(A; alg)
-            return V * sqrth_safe(D; kwargs...), invsqrth_safe(D; kwargs...) * V'
-        end
-    end
-end
-
-"""
-    gram_eigh_full(A::AbstractMatrix; alg=nothing, atol=0, rtol=eps(real(eltype(A)))^(3//4)) -> X
-
-Gram factorization of a Hermitian positive semi-definite matrix via its
-eigendecomposition (balanced eigh): returns `X = V * sqrth_safe(D; atol, rtol)`
-such that `A ≈ X * X'`, where `A = V * D * V'`. The square-root of `D` is
-absorbed symmetrically into the two factors of the eigendecomposition.
-Eigenvalues below `tol` (see [`pow_diag_safe`](@ref)) are clamped to zero.
-
-## Keyword arguments
-
-  - `alg`: forwarded to `MatrixAlgebraKit.eigh_full`.
-
-$(_clamp_kwargs_doc("A"))
-
-# Examples
-
-```jldoctest
-julia> using TensorAlgebra.MatrixAlgebra: gram_eigh_full
-
-julia> B = [1.0 0.5; 0.5 2.0];
-
-julia> A = B' * B;
-
-julia> X = gram_eigh_full(A);
-
-julia> X * X' ≈ A
-true
-```
-
-See also [`gram_eigh_full_with_pinv`](@ref).
-"""
-gram_eigh_full
-
-"""
-    gram_eigh_full_with_pinv(A::AbstractMatrix; alg=nothing, atol=0, rtol=eps(real(eltype(A)))^(3//4)) -> X, Y
-
-Like [`gram_eigh_full`](@ref), but additionally returns
-`Y = invsqrth_safe(D; atol, rtol) * V' ≈ pinv(X)`, a left inverse of `X`
-on the rank subspace: `Y * X ≈ I`. Eigenvalues below `tol` are clamped to
-zero in both factors.
-
-## Keyword arguments
-
-  - `alg`: forwarded to `MatrixAlgebraKit.eigh_full`.
-
-$(_clamp_kwargs_doc("A"))
-
-# Examples
-
-```jldoctest
-julia> using LinearAlgebra: I
-
-julia> using TensorAlgebra.MatrixAlgebra: gram_eigh_full_with_pinv
-
-julia> B = [1.0 0.5; 0.5 2.0];
-
-julia> A = B' * B;
-
-julia> X, Y = gram_eigh_full_with_pinv(A);
-
-julia> X * X' ≈ A
-true
-
-julia> Y * X ≈ I
-true
-```
-"""
-gram_eigh_full_with_pinv
-
 using MatrixAlgebraKit: MatrixAlgebraKit, TruncationStrategy
 
 struct TruncationDegenerate{Strategy <: TruncationStrategy, T <: Real} <: TruncationStrategy

src/TensorAlgebra.jl

@@ -1,15 +1,11 @@
 module TensorAlgebra
 
-export contract, contract!, dual, eig_full, eig_trunc, eig_vals, eigh_full, eigh_trunc,
-    eigh_vals, gram_eigh_full, gram_eigh_full_with_pinv, invsqrth_safe, isdual, left_null,
-    left_orth, left_polar, lq_compact, lq_full, project_hermitian, qr_compact,
-    qr_full, right_null, right_orth, right_polar, sqrth_invsqrth_safe, sqrth_safe,
-    svd_compact, svd_full, svd_trunc, svd_vals
+export contract, contract!, contractalign, dual, isdual, MatrixAlgebra
 
 if VERSION >= v"1.11.0-DEV.469"
     eval(
         Meta.parse(
-            "public biperm, bipartition, cat_similar, concatenate, concatenate!, ContractAlgorithm, contractopadd!, data, datatype, directsum, flattenlinear, label_type, matricizeopperm, permutedims, permutedims!, scalar, similar_map, TensorOperationsAlgorithm, to_range, tr, tryflattenlinear, ungrade, zero!, scale!, permuteddims, PermutedDims"
+            "public AbstractAlgorithm, add!, AddBroadcasted, addends, allocate_output, allocate_project, arguments, axes, bipartition, bipartition_axes, biperm, bipermutedims, bipermutedims!, bipermutedimsopadd!, cat_axis, cat_similar, check_input, concatenate, concatenate!, ConjBroadcasted, contractadd!, ContractAlgorithm, contractopadd!, contractperm, contractperm!, contractpermadd!, contractpermalign, contractpermopadd!, data, datatype, default_algorithm, dims2cat, directsum, eig_full, eig_trunc, eig_vals, eigh_full, eigh_trunc, eigh_vals, fill_map, flattenlinear, infer_aux_space, invsqrth_safe, is_output_view, is_projected, isidentitybiperm, label_type, left_null, left_orth, left_polar, LinearBroadcasted, linearbroadcasted, lq_compact, lq_full, matricize, MatricizeContract, matricizeop, matricizeop!, matricizeopcopy, matricizeopview, MatricizeStyle, MATRIX_FUNCTIONS, ndims, ndims_codomain, ndims_domain, one, ones_map, operation, output_axes, PermutedDims, permuteddims, permutedims, permutedims!, permutedimsadd!, permutedimsop, permutedimsopadd!, project, project!, project_aux, project_hermitian, projectto!, qr_compact, qr_full, rand_map, randn_map, right_null, right_orth, right_polar, scalar, scale!, ScaledBroadcasted, select_algorithm, similar_map, size, sqrth_invsqrth_safe, sqrth_safe, sum, svd_compact, svd_full, svd_trunc, svd_vals, TensorOperationsContract, to_range, tr, trivialrange, tryflattenlinear, tryproject, tryproject_aux, unchecked_project, unchecked_project_aux, ungrade, unmatricize, unmatricize!, unmatricize_factors, unproject, unscaled, zero!, zeros_map"
         )
     )
 end
@@ -25,6 +21,7 @@ include("concatenate.jl")
 include("directsum.jl")
 include("dual.jl")
 include("to_range.jl")
+include("algorithm.jl")
 include("contract/contractalgorithm.jl")
 include("contract/contract.jl")
 include("contract/contract_labels.jl")