[WIP] Rework the matricize, contract, and factorization surfaces for v0.21

Project.toml

@@ -1,13 +1,12 @@
 name = "TensorAlgebra"
 uuid = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
-version = "0.20.1"
+version = "0.21.0-DEV"
 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/TensorAlgebraTensorKitExt.jl

@@ -243,44 +243,41 @@ 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 a `TensorMap`
+# over the regrouped space and the write is the ordinary permuted-add. `bipermutedimsopadd!` above
+# routes that through `tensoradd!`, which realizes the permutation, the `op === conj` conjugation
+# and the scaling in one call, so no separate handling of `op` is needed here.
+function TensorAlgebra.allocate_output(
+        ::typeof(TensorAlgebra.matricizeop), ::TensorKitMatricize, op,
+        t::AbstractTensorMap, perm_codomain, perm_domain
     )
+    return similar(t, permute(space(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.
 function TensorAlgebra.unmatricize(
         ::TensorKitMatricize, m::AbstractTensorMap, axes_codomain, axes_domain
     )
@@ -296,7 +293,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

ext/TensorAlgebraTensorOperationsExt/TensorAlgebraTensorOperationsExt.jl

@@ -21,42 +21,7 @@ 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!(
+function TA.contractpermopadd!(
         algorithm::TensorOperationsAlgorithm,
         a_dest, perm_dest_codomain, perm_dest_domain,
         op1, a1, perm1_codomain, perm1_domain,
@@ -89,7 +54,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,15 +1,12 @@
 module MatrixAlgebra
 
-export gram_eigh_full,
-    gram_eigh_full_with_pinv,
-    invsqrt_diag_safe,
+export invsqrt_diag_safe,
     invsqrth_safe,
     pow_diag_safe,
     pow_diag_safe!,
     powh_safe,
     sqrt_diag_safe,
-    sqrth_safe,
-    sqrth_invsqrth_safe
+    sqrth_safe
 
 using LinearAlgebra: LinearAlgebra, Diagonal, isdiag, norm
 using MatrixAlgebraKit: MatrixAlgebraKit as MAK
@@ -165,120 +162,6 @@ $(_clamp_kwargs_doc("M"))
 """
 invsqrth_safe(M; kwargs...) = powh_safe(M, -1 // 2; kwargs...)
 
-"""
-    sqrth_invsqrth_safe(M; alg=nothing, atol=0, rtol=eps(real(eltype(M)))^(3//4)) -> M^(1//2), M^(-1//2)
-
-Square root and pseudo-inverse square root of a Hermitian positive
-semi-definite matrix, from a single eigendecomposition. Equivalent
-to `(sqrth_safe(M; ...), invsqrth_safe(M; ...))` but with the
-eigendecomposition computed once. Eigenvalues below tolerance are clamped
-to zero in both factors (Moore-Penrose convention for the inverse).
-
-The input must be Hermitian (as for `MatrixAlgebraKit.eigh_full`): project
-with `MatrixAlgebraKit.project_hermitian` first if it is Hermitian only up
-to numerical noise.
-
-## Keyword arguments
-
-  - `alg`: forwarded to `MatrixAlgebraKit.eigh_full`.
-
-$(_clamp_kwargs_doc("M"))
-"""
-function sqrth_invsqrth_safe(M; alg = nothing, kwargs...)
-    if isdiag(M)
-        return pow_diag_safe(M, 1 // 2; kwargs...), pow_diag_safe(M, -1 // 2; kwargs...)
-    end
-    D, V = MAK.eigh_full(M; alg)
-    return V * pow_diag_safe(D, 1 // 2; kwargs...) * V',
-        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,16 @@
 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,
+export contract, contract!, contractalign, dual, eig_full, eig_trunc, eig_vals, eigh_full,
+    eigh_trunc,
+    eigh_vals, 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,
+    qr_full, right_null, right_orth, right_polar, sqrth_safe,
     svd_compact, svd_full, svd_trunc, svd_vals
 
 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 allocate_output, biperm, bipartition, cat_similar, check_input, concatenate, concatenate!, ContractAlgorithm, contractadd!, contractopadd!, contractperm, contractperm!, contractpermadd!, contractpermopadd!, data, datatype, directsum, flattenlinear, is_output_view, label_type, matricize, matricizeop, matricizeop!, matricizeopcopy, matricizeopview, output_axes, select_algorithm, default_algorithm, permutedims, permutedims!, scalar, similar_map, TensorOperationsAlgorithm, to_range, tr, tryflattenlinear, ungrade, zero!, scale!, permuteddims, PermutedDims"
         )
     )
 end
@@ -25,6 +26,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")

src/algorithm.jl

@@ -0,0 +1,39 @@
+"""
+    TensorAlgebra.select_algorithm(f, args...; alg = nothing, kwargs...)
+
+Resolve the algorithm operation `f` should run with on `args`. An `alg` of `nothing` defers to
+[`TensorAlgebra.default_algorithm`](@ref); anything else is validated and passed through.
+
+This is the forward-facing layer the per-operation resolvers sit under, so a caller that does not
+care which operation it is dispatching writes `select_algorithm(f, ...)` and a backend registers
+its choice on `default_algorithm(f, ...)`.
+"""
+function select_algorithm(f, args...; alg = nothing, kwargs...)
+    isnothing(alg) && return default_algorithm(f, args...; kwargs...)
+    return select_algorithm_specified(f, alg, args...; kwargs...)
+end
+
+"""
+    TensorAlgebra.default_algorithm(f, args...)
+    TensorAlgebra.default_algorithm(f, argtypes::Type...)
+
+The algorithm operation `f` runs with on `args` when the caller names none. The types form is the
+registration point for a storage type; the values form defaults to it.
+
+Each operation bridges to its own resolver, so `default_algorithm(contract, A1, A2)` is
+[`TensorAlgebra.default_contract_algorithm`](@ref).
+"""
+function default_algorithm(f, args...; kwargs...)
+    return default_algorithm(f, map(typeof, args)...; kwargs...)
+end
+function default_algorithm(f, argtypes::Type...; kwargs...)
+    return throw(MethodError(default_algorithm, (f, argtypes...)))
+end
+
+# `alg` named something. A resolved algorithm object passes through; anything else is a caller
+# error, reported against the operation rather than as a `MethodError` from inside the resolver.
+function select_algorithm_specified(f, alg, args...; kwargs...)
+    return throw(
+        ArgumentError("`$alg` is not an algorithm for `$f`")
+    )
+end