Operand-driven diagonal algebra for the dense path

Project.toml

@@ -1,6 +1,6 @@
 name = "TensorAlgebra"
 uuid = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
-version = "0.20.0"
+version = "0.20.0-DEV"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
 
 [workspace]

ext/TensorAlgebraTensorKitExt.jl

@@ -66,10 +66,10 @@ TensorAlgebra.sum(t::AbstractTensorMap; kwargs...) = Base.sum(convert(Array, t);
 # exactly what TensorKit's `similar(t, T, codomain, domain)` wants, so build the two
 # `ProductSpace`s directly.
 function TensorAlgebra.similar_map(
-        a::AbstractTensorMap, ::Type{T}, codomain_axes, domain_axes
+        a::AbstractTensorMap, ::Type{T}, axes_codomain, axes_domain
     ) where {T}
     S = spacetype(a)
-    return similar(a, T, ProductSpace{S}(codomain_axes...), ProductSpace{S}(domain_axes...))
+    return similar(a, T, ProductSpace{S}(axes_codomain...), ProductSpace{S}(axes_domain...))
 end
 
 # A plain-array prototype with native (space) axes is the operator/state construction case: `raw`
@@ -81,22 +81,22 @@ end
 # entries below can read it from whichever of codomain/domain is non-empty and share one builder,
 # mirroring `_map_homspace` and the map constructors.
 function similar_tensormap(
-        raw::AbstractArray, ::Type{T}, ::Type{S}, codomain_axes, domain_axes
+        raw::AbstractArray, ::Type{T}, ::Type{S}, axes_codomain, axes_domain
     ) where {T, S <: ElementarySpace}
     A = Base.promote_op(similar, typeof(raw), Type{T}, Int)
-    return TensorMapWithStorage{T, A}(undef, _map_homspace(S, codomain_axes, domain_axes))
+    return TensorMapWithStorage{T, A}(undef, _map_homspace(S, axes_codomain, axes_domain))
 end
 function TensorAlgebra.similar_map(
         raw::AbstractArray, ::Type{T},
-        codomain_axes::Tuple{S, Vararg{S}}, domain_axes::Tuple{Vararg{S}}
+        axes_codomain::Tuple{S, Vararg{S}}, axes_domain::Tuple{Vararg{S}}
     ) where {T, S <: ElementarySpace}
-    return similar_tensormap(raw, T, S, codomain_axes, domain_axes)
+    return similar_tensormap(raw, T, S, axes_codomain, axes_domain)
 end
 function TensorAlgebra.similar_map(
         raw::AbstractArray, ::Type{T},
-        codomain_axes::Tuple{}, domain_axes::Tuple{S, Vararg{S}}
+        axes_codomain::Tuple{}, axes_domain::Tuple{S, Vararg{S}}
     ) where {T, S <: ElementarySpace}
-    return similar_tensormap(raw, T, S, codomain_axes, domain_axes)
+    return similar_tensormap(raw, T, S, axes_codomain, axes_domain)
 end
 
 # ===============================  zeros_map / randn_map / rand_map  ========================
@@ -107,32 +107,32 @@ end
 # `S` is passed to `_map_homspace` explicitly so the two dispatch entries per constructor can
 # read it from whichever of the codomain/domain is non-empty and share one builder; an empty
 # axis tuple gives the unit space `ProductSpace{S}()`.
-function _map_homspace(::Type{S}, codomain_axes, domain_axes) where {S <: ElementarySpace}
-    return ProductSpace{S}(codomain_axes...) ← ProductSpace{S}(domain_axes...)
+function _map_homspace(::Type{S}, axes_codomain, axes_domain) where {S <: ElementarySpace}
+    return ProductSpace{S}(axes_codomain...) ← ProductSpace{S}(axes_domain...)
 end
 function TensorAlgebra.zeros_map(
-        ::Type{T}, codomain_axes::Tuple{S, Vararg{S}}, domain_axes::Tuple{Vararg{S}}
+        ::Type{T}, axes_codomain::Tuple{S, Vararg{S}}, axes_domain::Tuple{Vararg{S}}
     ) where {T, S <: ElementarySpace}
-    return TensorKit.zeros(T, _map_homspace(S, codomain_axes, domain_axes))
+    return TensorKit.zeros(T, _map_homspace(S, axes_codomain, axes_domain))
 end
 function TensorAlgebra.zeros_map(
-        ::Type{T}, codomain_axes::Tuple{}, domain_axes::Tuple{S, Vararg{S}}
+        ::Type{T}, axes_codomain::Tuple{}, axes_domain::Tuple{S, Vararg{S}}
     ) where {T, S <: ElementarySpace}
-    return TensorKit.zeros(T, _map_homspace(S, codomain_axes, domain_axes))
+    return TensorKit.zeros(T, _map_homspace(S, axes_codomain, axes_domain))
 end
 for (f, g) in ((:randn_map, :randn), (:rand_map, :rand))
     @eval begin
         function TensorAlgebra.$f(
                 rng::AbstractRNG, ::Type{T},
-                codomain_axes::Tuple{S, Vararg{S}}, domain_axes::Tuple{Vararg{S}}
+                axes_codomain::Tuple{S, Vararg{S}}, axes_domain::Tuple{Vararg{S}}
             ) where {T, S <: ElementarySpace}
-            return TensorKit.$g(rng, T, _map_homspace(S, codomain_axes, domain_axes))
+            return TensorKit.$g(rng, T, _map_homspace(S, axes_codomain, axes_domain))
         end
         function TensorAlgebra.$f(
                 rng::AbstractRNG, ::Type{T},
-                codomain_axes::Tuple{}, domain_axes::Tuple{S, Vararg{S}}
+                axes_codomain::Tuple{}, axes_domain::Tuple{S, Vararg{S}}
             ) where {T, S <: ElementarySpace}
-            return TensorKit.$g(rng, T, _map_homspace(S, codomain_axes, domain_axes))
+            return TensorKit.$g(rng, T, _map_homspace(S, axes_codomain, axes_domain))
         end
     end
 end
@@ -142,7 +142,7 @@ end
 # `TensorMap` is not an `AbstractArray`, so the generic `copyto!` default does not apply; delegate
 # to TensorKit's `project_symmetric!`, which fills the symmetry-allowed blocks from the dense data
 # and discards any component outside the block structure. Composed with the map constructors above,
-# this makes `project(dense, codomain_axes, domain_axes)` build a `TensorMap` from a dense matrix.
+# this makes `project(dense, axes_codomain, axes_domain)` build a `TensorMap` from a dense matrix.
 # `project_symmetric!` requires a matching dense size, so reshape `src` to `size(dest)` first (a
 # no-op when the ranks already match); this lets a lower-rank `src` omit trailing length-1 axes,
 # matching the generic `projectto!`, and rejects a genuine shape mismatch.
@@ -168,14 +168,14 @@ end
 # whichever side is non-empty, the same two-entry split `similar_map` uses. `project`'s allocation
 # stays generic (strict `similar_map`); only the `project_aux` derivation is `TensorMap`-specific.
 function TensorAlgebra.infer_aux_space(
-        raw::AbstractArray, codomain_axes::Tuple{S, Vararg{S}}, domain_axes::Tuple{Vararg{S}}
+        raw::AbstractArray, axes_codomain::Tuple{S, Vararg{S}}, axes_domain::Tuple{Vararg{S}}
     ) where {S <: ElementarySpace}
-    return infer_aux_space_tensormap(raw, S, codomain_axes, domain_axes)
+    return infer_aux_space_tensormap(raw, S, axes_codomain, axes_domain)
 end
 function TensorAlgebra.infer_aux_space(
-        raw::AbstractArray, codomain_axes::Tuple{}, domain_axes::Tuple{S, Vararg{S}}
+        raw::AbstractArray, axes_codomain::Tuple{}, axes_domain::Tuple{S, Vararg{S}}
     ) where {S <: ElementarySpace}
-    return infer_aux_space_tensormap(raw, S, codomain_axes, domain_axes)
+    return infer_aux_space_tensormap(raw, S, axes_codomain, axes_domain)
 end
 
 # The space of `raw`'s trailing auxiliary axis, derived so the projected result is
@@ -184,11 +184,11 @@ end
 # follows, so the aux slices must appear in that order. The result may span several sectors (a
 # direct-sum, MPO-style virtual leg).
 function infer_aux_space_tensormap(
-        raw, ::Type{S}, codomain_axes, domain_axes
+        raw, ::Type{S}, axes_codomain, axes_domain
     ) where {S <: ElementarySpace}
-    aux_dim = length(codomain_axes) + length(domain_axes) + 1
+    aux_dim = length(axes_codomain) + length(axes_domain) + 1
     aux_length = size(raw, aux_dim)
-    content = fuse(codomain_axes..., dual.(domain_axes)...)
+    content = fuse(axes_codomain..., dual.(axes_domain)...)
     # Probe the surplus axis slice by slice: a slice keeps the aux axis (width `dim(s)`), so its
     # rank matches the candidate axes exactly and `tryproject` allocates, fills, and round-trip-
     # verifies without re-entering the derivation branch. This builds one `TensorMap` per candidate
@@ -197,7 +197,7 @@ function infer_aux_space_tensormap(
     function slice_is_covariant(r, s)
         slice = selectdim(raw, aux_dim, r)
         return !isnothing(
-            TensorAlgebra.tryproject(slice, codomain_axes, (domain_axes..., S(s => 1)))
+            TensorAlgebra.tryproject(slice, axes_codomain, (axes_domain..., S(s => 1)))
         )
     end
     seccounts = Pair{TensorKit.sectortype(S), Int}[]
@@ -264,10 +264,10 @@ end
 # codomain-facing (un-dualized), which is exactly TensorKit's domain convention, so they build the
 # domain `ProductSpace` directly.
 function TensorAlgebra.unmatricize(
-        ::TensorKitMatricize, m::AbstractTensorMap, codomain_axes, domain_axes
+        ::TensorKitMatricize, m::AbstractTensorMap, axes_codomain, axes_domain
     )
     S = spacetype(m)
-    dest = ProductSpace{S}(codomain_axes...) ← ProductSpace{S}(domain_axes...)
+    dest = ProductSpace{S}(axes_codomain...) ← ProductSpace{S}(axes_domain...)
     space(m) == dest ||
         throw(ArgumentError("`unmatricize` space `$dest` does not match `$(space(m))`"))
     return m

src/TensorAlgebra.jl

@@ -23,14 +23,14 @@ include("permutedimsadd.jl")
 include("matricize.jl")
 include("concatenate.jl")
 include("directsum.jl")
-include("diagonal.jl")
 include("dual.jl")
 include("to_range.jl")
 include("contract/contractalgorithm.jl")
 include("contract/contract.jl")
 include("contract/contract_labels.jl")
 include("contract/biperms.jl")
 include("contract/allocate_output.jl")
+include("diagonal.jl")
 include("contract/contract_matricize.jl")
 include("factorizations.jl")
 include("matrixfunctions.jl")

src/contract/allocate_output.jl

@@ -64,7 +64,13 @@ function output_axes(
     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)
+    axes_codomain_dest, axes_domain_dest = bipartition(
+        axes_uncontracted, perm_dest_codomain, perm_dest_domain
+    )
+    # The operand axes are stored/dualized, so un-dualize the domain axes into the codomain-facing
+    # construction convention shared by `allocate_contract_output`, `similar_map`, and `unmatricize`
+    # (a no-op on dense axes).
+    return axes_codomain_dest, conj.(axes_domain_dest)
 end
 
 # TODO: Use `ArrayLayouts`-like `MulAdd` object,
@@ -84,14 +90,21 @@ function allocate_output(
         perm2_codomain,
         perm2_domain
     )
-    codomain_axes_dest, domain_axes_dest = output_axes(
+    axes_codomain_dest, axes_domain_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))
-    # `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)))
+    T = Base.promote_op(matprod, eltype(a1), eltype(a2))
+    return allocate_contract_output(a1, a2, T, axes_codomain_dest, axes_domain_dest)
+end
+
+# Allocate the output container for `contract`: the operand types, the output element type and
+# axes (domain codomain-facing), and the output's codomain/domain leg counts (the axes tuple
+# lengths) select the container type. Internal to TensorAlgebra, not a public extension point:
+# the leg counts identify the contraction pattern only for matrix-shaped operands (see the
+# `Diagonal` method in `diagonal.jl`), so external structured types should not overload it.
+function allocate_contract_output(a1, a2, T, axes_codomain::Tuple, axes_domain::Tuple)
+    return zero!(similar_map(a1, T, axes_codomain, axes_domain))
 end

src/diagonal.jl

@@ -5,10 +5,9 @@ using LinearAlgebra: Diagonal
 # an operation is still diagonal. These methods hook the lowest-level primitives, so the
 # convenience wrappers built on them (`bipermutedims`, `permutedimsadd!`, `add!`, and the
 # matrix functions, which all route through `bipermutedimsopadd!` and `allocate_output`)
-# preserve `Diagonal`. Structure is given up (via the generic reshape path) only where the
-# result genuinely is not diagonal: vectorizing matricizations and `Diagonal`/dense mixing
-# (the latter falls back to Base's dense `similar`, since `contract` allocates from a flat
-# axis tuple, not a `BiTuple`).
+# preserve `Diagonal`. Structure is given up only where the result genuinely is not diagonal:
+# vectorizing matricizations, a non-`{1,1}` bond-split `unmatricize`, `Diagonal`/dense mixing,
+# and `contract` patterns other than the single-contracted-leg matmul (`Diagonal * Diagonal`).
 
 # Permuting the two axes of a square `Diagonal` (identity or transpose) leaves it
 # unchanged, so the lazy permutation is the matrix itself.
@@ -45,9 +44,33 @@ end
 # A `Diagonal` is already a matrix; the `(1 codomain, 1 domain)` matricization is the identity
 # reshape, so return it directly (maybe-alias, matching `matricize`'s general contract).
 matricize(::ReshapeMatricize, a::Diagonal, ::Val{1}) = a
+# A `{1,1}` unmatricize (one codomain axis, one domain axis) is the endomorphism identity: the
+# result stays `Diagonal`, so return `m` directly. The generic `check_input(unmatricize, ...)`
+# validates the axis lengths against `m`'s size.
 function unmatricize(
         ::ReshapeMatricize, m::Diagonal,
-        ::Tuple{<:AbstractUnitRange}, ::Tuple{<:AbstractUnitRange}
+        axes_codomain::Tuple{<:AbstractUnitRange}, axes_domain::Tuple{<:AbstractUnitRange}
     )
+    check_input(unmatricize, m, axes_codomain, axes_domain)
     return m
 end
+# Any other split is a genuine bond-split (for example `unmatricize(D[4×4], (2, 2), (4,))`) whose
+# result is not representable as a `Diagonal`, so densify and reshape like a dense matrix.
+# `copyto!(similar(m, axes(m)), m)` densifies while preserving `m`'s array backend (a plain
+# `Array` would force the result onto the CPU).
+function unmatricize(
+        style::ReshapeMatricize, m::Diagonal, axes_codomain::Tuple, axes_domain::Tuple
+    )
+    return unmatricize(style, copyto!(similar(m, axes(m)), m), axes_codomain, axes_domain)
+end
+
+# Contracting two `Diagonal`s over a single leg is the matmul/endomorphism pattern
+# `Diagonal * Diagonal = Diagonal` (all transpose variants `[i,j]*[j,k]`, `[i,j]*[k,j]`, ...),
+# whose `{1,1}` output stays `Diagonal`, so allocate one. Every other output shape (rank-4 outer
+# product, scalar full contraction) is not representable as a `Diagonal` and falls back to the
+# generic dense allocation, matching `Diagonal`/dense mixing.
+function allocate_contract_output(
+        a1::Diagonal, a2::Diagonal, T, axes_codomain::Tuple{Any}, axes_domain::Tuple{Any}
+    )
+    return Diagonal(zero!(similar(a1.diag, T, (only(axes_codomain),))))
+end

src/factorizations.jl

@@ -1,8 +1,9 @@
 using LinearAlgebra: LinearAlgebra
 using MatrixAlgebraKit: MatrixAlgebraKit
 
-# Each factorization reconstructs its factors with `unmatricize`, reading the freshly created
-# bond axis off the factor itself: it is the factor's last axis on a codomain factor
+# Each factorization reconstructs its outer factors with `unmatricize`, while a spectrum factor
+# (`S`, `D`) is returned bare at the matrix level. The freshly created bond axis is read off the
+# factor itself: it is the factor's last axis on a codomain factor
 # (`[group…, bond]`) and its first axis on a domain factor (`[bond, group…]`), on every backend
 # (a fusing backend returns a rank-2 factor, a `TensorMap` keeps the group's original legs). The
 # bond is dualized to codomain-facing form (`conj`, a no-op on a dense axis) when it lands on the
@@ -261,7 +262,7 @@ for f in (:svd_compact, :svd_full)
             U, S, Vᴴ = MatrixAlgebraKit.$f(A_mat; kwargs...)
             axes_codomain, axes_domain = bipartition_axes(axes(A), ndims_codomain)
             return unmatricize(style, U, axes_codomain, (conj(axes(U, ndims(U))),)),
-                unmatricize(style, S, (axes(S, 1),), (conj(axes(S, 2)),)),
+                S,
                 unmatricize(style, Vᴴ, (axes(Vᴴ, 1),), axes_domain)
         end
         function $f(A, ndims_codomain::Val; kwargs...)
@@ -278,7 +279,7 @@ function svd_trunc(style::MatricizeStyle, A, ndims_codomain::Val; kwargs...)
     U, S, Vᴴ, ϵ = MatrixAlgebraKit.svd_trunc(A_mat; kwargs...)
     axes_codomain, axes_domain = bipartition_axes(axes(A), ndims_codomain)
     return unmatricize(style, U, axes_codomain, (conj(axes(U, ndims(U))),)),
-        unmatricize(style, S, (axes(S, 1),), (conj(axes(S, 2)),)),
+        S,
         unmatricize(style, Vᴴ, (axes(Vᴴ, 1),), axes_domain),
         ϵ
 end
@@ -287,14 +288,15 @@ function svd_trunc(A, ndims_codomain::Val; kwargs...)
 end
 
 # Eigendecomposition: `D` is the rank × rank spectrum and `V` carries the codomain axes plus a
-# trailing rank axis. Both are unmatricized back to the array type, as in `svd_*`.
+# trailing rank axis. `D` is returned bare (its axis is the internal bond, so there is nothing to
+# unfold); `V` is unmatricized back to the array type, as in `svd_*`.
 for f in (:eigh_full, :eig_full, :eigh_trunc, :eig_trunc)
     @eval begin
         function $f(style::MatricizeStyle, A, ndims_codomain::Val; kwargs...)
             A_mat = matricize(style, A, ndims_codomain)
             D, V = MatrixAlgebraKit.$f(A_mat; kwargs...)
             axes_codomain = first(bipartition(axes(A), ndims_codomain))
-            return unmatricize(style, D, (axes(D, 1),), (conj(axes(D, 2)),)),
+            return D,
                 unmatricize(style, V, axes_codomain, (conj(axes(V, ndims(V))),))
         end
         function $f(A, ndims_codomain::Val; kwargs...)
@@ -812,8 +814,8 @@ one
 function one!!(style::MatricizeStyle, A, ndims_codomain::Val; kwargs...)
     A_mat = matricize(style, A, ndims_codomain)
     MatrixAlgebraKit.one!(A_mat)
-    codomain_axes, domain_axes = bipartition_axes(axes(A), ndims_codomain)
-    return unmatricize(style, A_mat, codomain_axes, domain_axes)
+    axes_codomain, axes_domain = bipartition_axes(axes(A), ndims_codomain)
+    return unmatricize(style, A_mat, axes_codomain, axes_domain)
 end
 function one!!(A, ndims_codomain::Val; kwargs...)
     return one!!(MatricizeStyle(A), A, ndims_codomain; kwargs...)

src/linearbroadcasted.jl

@@ -145,7 +145,9 @@ arguments(a::AddBroadcasted) = addends(a)
 # Mul — lazy matrix multiplication (standalone, not LinearBroadcasted)
 # ---------------------------------------------------------------------------- #
 
-# Same as `LinearAlgebra.matprod`, but duplicated here since it is private.
+# The element type of a matrix-product accumulation, modeled on `LinearAlgebra.matprod` (private
+# there, so inlined): a sum of products, so `promote_op(matprod, ...)` widens exactly as an
+# accumulating inner loop would.
 matprod(x, y) = x * y + x * y
 
 struct Mul{A, B}