Move to GradedArrays 0.16 and TensorAlgebra 0.20, simplify named broadcasting

Project.toml

@@ -1,6 +1,6 @@
 name = "ITensorBase"
 uuid = "4795dd04-0d67-49bb-8f44-b89c448a1dc7"
-version = "0.13.13"
+version = "0.14.0-DEV"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
 
 [workspace]
@@ -31,6 +31,10 @@ OMEinsumContractionOrders = "6f22d1fd-8eed-4bb7-9776-e7d684900715"
 TensorKit = "07d1fe3e-3e46-537d-9eac-e9e13d0d4cec"
 TensorKitSectors = "13a9c161-d5da-41f0-bcbd-e1a08ae0647f"
 
+[sources.TensorAlgebra]
+rev = "main"
+url = "https://github.com/ITensor/TensorAlgebra.jl"
+
 [extensions]
 ITensorBaseAdaptExt = "Adapt"
 ITensorBaseGradedArraysExt = ["GradedArrays", "TensorKitSectors"]
@@ -45,14 +49,14 @@ Adapt = "4.1.1"
 ArrayLayouts = "1.11"
 Combinatorics = "1"
 ConstructionBase = "1.6"
-GradedArrays = "0.15"
+GradedArrays = "0.16"
 LinearAlgebra = "1.10"
 MatrixAlgebraKit = "0.2, 0.3, 0.4, 0.5, 0.6"
 Mooncake = "0.4.202, 0.5"
 OMEinsumContractionOrders = "1.3"
 Random = "1.10"
 SimpleTraits = "0.9.4"
-TensorAlgebra = "0.18"
+TensorAlgebra = "0.20"
 TensorKit = "0.17"
 TensorKitSectors = "0.3.9"
 TermInterface = "2"

docs/Project.toml

@@ -12,9 +12,9 @@ path = ".."
 
 [compat]
 Documenter = "1"
-ITensorBase = "0.13"
+ITensorBase = "0.14"
 ITensorFormatter = "0.2.27"
 Literate = "2"
 MatrixAlgebraKit = "0.2, 0.3, 0.4, 0.5, 0.6"
-TensorAlgebra = "0.18"
+TensorAlgebra = "0.20"
 Test = "1.10"

examples/Project.toml

@@ -6,5 +6,5 @@ MatrixAlgebraKit = "6c742aac-3347-4629-af66-fc926824e5e4"
 path = ".."
 
 [compat]
-ITensorBase = "0.13"
+ITensorBase = "0.14"
 MatrixAlgebraKit = "0.2, 0.3, 0.4, 0.5, 0.6"

src/broadcast.jl

@@ -21,14 +21,43 @@ end
 # `AbstractArray`); without this the default `broadcastable` wraps it in a `Ref`.
 BC.broadcastable(a::AbstractNamedTensor) = a
 
-broadcasted_unnamed(x::Number, names) = x
-function broadcasted_unnamed(a::AbstractNamedTensor, names)
-    # An operand already aligned to the destination names (the first operand always, and the
-    # common case for the rest) needs no permutation, avoiding a `getperm` allocation and the
-    # identity `permuteddims` wrapper. Skipping it makes a small add several times slower.
+# Unname a flattened named `LinearBroadcasted`. A single scaled/conjugated operand is already aligned to
+# the output names, so no permutation is needed and its codomain/domain split is kept; only a sum needs
+# its addends aligned. (Flattening distributes scaling/conjugation over `+`, so a `Scaled`/`Conj` node
+# never wraps an `Add`, and the no-permutation recursion below never reaches one.)
+unnamed_linear(a::TA.LinearBroadcasted, names) = unnamed_linear(a)
+unnamed_linear(a::TA.AddBroadcasted, names) = unnamed_linear_aligned(a, names)
+
+# No permutation: strip names down the expression tree via the `operation`/`arguments` term interface.
+function unnamed_linear(a::TA.LinearBroadcasted)
+    return TA.linearbroadcasted(TA.operation(a), map(unnamed_linear, TA.arguments(a))...)
+end
+unnamed_linear(a::AbstractNamedTensor) = unnamed(a)
+unnamed_linear(a::Number) = a
+
+# Align every leaf to `names` through the `PermutedDims` wrapper (all-codomain output). Used for a sum's
+# addends and for every in-place `copyto!` (aligned to the destination).
+function unnamed_linear_aligned(a::TA.LinearBroadcasted, names)
+    return TA.linearbroadcasted(
+        TA.operation(a), map(x -> unnamed_linear_aligned(x, names), TA.arguments(a))...
+    )
+end
+function unnamed_linear_aligned(a::AbstractNamedTensor, names)
+    return _broadcast_permuteddims(unnamed(a), getperm(dimnames(a), names))
+end
+unnamed_linear_aligned(a::Number, names) = a
+
+# Non-linear fallback: unname a general `Broadcasted` by aligning each operand to `names`, so Base's
+# generic broadcast can run (all-codomain output). Only the linear path preserves the split.
+unnamed_broadcasted(x::Number, names) = x
+function unnamed_broadcasted(a::AbstractNamedTensor, names)
+    # An operand already aligned to `names` needs no permutation, skipping the identity wrapper.
     dimnames(a) == names && return unnamed(a)
     return _broadcast_permuteddims(unnamed(a), getperm(dimnames(a), names))
 end
+function unnamed_broadcasted(bc::Broadcasted, names)
+    return broadcasted(bc.f, Base.Fix2(unnamed_broadcasted, names).(bc.args)...)
+end
 # Broadcasting-only alignment: unlike the public `unnamed(a, names)` (which returns a
 # `Base.PermutedDimsArray`, a full array), this wraps in `TensorAlgebra.PermutedDims`, which stores
 # the permutation in a field rather than a type parameter, so it builds cheaply and type-stably
@@ -41,20 +70,8 @@ end
 @noinline function _broadcast_permuteddims(array, perm)
     return TA.PermutedDims(array, ntuple(i -> perm[i], Val(TA.ndims(array))))
 end
-function broadcasted_unnamed(bc::Broadcasted, names)
-    return broadcasted(bc.f, Base.Fix2(broadcasted_unnamed, names).(bc.args)...)
-end
-
-# A bare (unnamed) array operand, used as an allocation prototype so a broadcast
-# result inherits the operands' backend (e.g. graded) rather than a lazy permuted
-# wrapper's `similar` (which can drop the backend).
-unnamed_prototype(bc::Broadcasted) = unnamed_prototype(bc.args...)
-unnamed_prototype(arg::AbstractNamedTensor, args...) = unnamed(arg)
-unnamed_prototype(arg::Broadcasted, args...) = unnamed_prototype(arg.args..., args...)
-unnamed_prototype(arg, args...) = unnamed_prototype(args...)
-
 # Skip Base's shape-combination step: named broadcasts don't need the `NamedUnitRange` axis
-# machinery. Name compatibility is handled by the per-operand alignment in `broadcasted_unnamed`
+# machinery. Name compatibility is handled by the per-operand alignment in `unnamed_linear_aligned`
 # (via `getperm`), and unnamed-shape compatibility by TensorAlgebra.
 BC.instantiate(bc::Broadcasted{<:AbstractNamedTensorStyle}) = bc
 
@@ -65,38 +82,21 @@ _dimnames(bc::Broadcasted, args...) = _dimnames(bc.args..., args...)
 _dimnames(_, args...) = _dimnames(args...)
 dimnames(bc::Broadcasted) = _dimnames(bc.args...)
 
-# The result element type of a linear combination, from the concrete unnamed leaves at runtime.
-# `eltype(::LinearBroadcasted)` uses `Base.promote_op`, which runs a live inference call here
-# because the leaves wrap a named tensor's (non-inferrable) backing array, so promote the
-# concrete `eltype`s instead.
-_lineareltype(a::AbstractArray) = eltype(a)
-function _lineareltype(s::TA.ScaledBroadcasted)
-    return promote_type(typeof(TA.coeff(s)), _lineareltype(TA.unscaled(s)))
-end
-_lineareltype(s::TA.AddBroadcasted) = promote_type(map(_lineareltype, TA.addends(s))...)
-
 function Base.copy(bc::Broadcasted{<:AbstractNamedTensorStyle})
     nms = dimnames(bc)
-    dest_unnamed = _copy_unnamed(broadcasted_unnamed(bc, nms), unnamed_prototype(bc))
-    return nameddims(dest_unnamed, nms)
+    return nameddims(_copy_unnamed(bc, nms), nms)
 end
 
-# Function barrier: `broadcasted_unnamed` and `unnamed_prototype` produce concretely-typed
-# values whose *inferred* types are abstract (the named backing array is abstract), so this
-# call re-specializes on the concrete runtime types and everything below is type-stable
-# (`eltype(lb)` is now inferrable, no runtime `promote_op`). Inlining the body into `copy`
-# instead costs one extra allocation per call.
-#
-# Allocate from `axes(lb)`, the flattened expression's axes, rather than the prototype's own:
-# an axis-changing operand (a `conj` leaf dualizes its axes) makes them differ, and the
-# destination must match the expression. All axes go in the codomain (empty domain), the
-# all-codomain output convention `@tensor` uses for an unbipartitioned left-hand side; on a
-# non-bipartitioned backend (a dense array) `similar_map` with an empty domain is a plain
-# `similar` over `axes(lb)`.
-function _copy_unnamed(bc_unnamed, prototype)
-    lb = TA.tryflattenlinear(bc_unnamed)
-    isnothing(lb) && return copy(bc_unnamed)
-    return copyto!(TA.similar_map(prototype, eltype(lb), axes(lb), ()), lb)
+# Function barrier: `bc`'s named leaves are abstractly typed, so re-dispatching on the concrete `bc`
+# here keeps the flatten/unname/materialize below type-stable. A linear expression folds to a
+# `LinearBroadcasted` and materializes through `copy(lb)`, whose allocation (`similar(lb)`) is the
+# unnamed backend's own broadcast-style `similar`, so the result inherits the backend (dense, graded,
+# ...) and `unnamed_linear` keeps a single scaled/conjugated operand's codomain/domain split. A
+# non-linear expression falls back to unnaming the raw `Broadcasted` and Base's generic broadcast.
+@noinline function _copy_unnamed(bc, nms)
+    lb = TA.tryflattenlinear(bc)
+    isnothing(lb) && return copy(unnamed_broadcasted(bc, nms))
+    return copy(unnamed_linear(lb, nms))
 end
 
 # `Base.Broadcast.materialize!` otherwise reconstructs the broadcast over `axes(dest)` and
@@ -115,17 +115,16 @@ function Base.copyto!(
         dest::AbstractNamedTensor,
         bc::Broadcasted{<:AbstractNamedTensorStyle}
     )
-    _copyto_unnamed!(unnamed(dest), broadcasted_unnamed(bc, dimnames(dest)))
+    _copyto_unnamed!(unnamed(dest), bc, dimnames(dest))
     return dest
 end
 
-# Function barrier mirroring `_copy_unnamed`: `unnamed(dest)` and `broadcasted_unnamed`
-# have abstract inferred types (the named backing array is abstract), so this call
-# re-specializes on the concrete runtime types and the flatten/lower below is type-stable.
-function _copyto_unnamed!(dest_unnamed, bc_unnamed)
-    lb = TA.tryflattenlinear(bc_unnamed)
-    isnothing(lb) && return copyto!(dest_unnamed, bc_unnamed)
-    return copyto!(dest_unnamed, lb)
+# Function barrier mirroring `_copy_unnamed`. In place, so every operand aligns to `dest`; non-linear
+# falls back to Base's generic in-place broadcast.
+@noinline function _copyto_unnamed!(dest_unnamed, bc, nms)
+    lb = TA.tryflattenlinear(bc)
+    isnothing(lb) && return copyto!(dest_unnamed, unnamed_broadcasted(bc, nms))
+    return copyto!(dest_unnamed, unnamed_linear_aligned(lb, nms))
 end
 
 # Operator-preserving broadcasting.
@@ -169,8 +168,7 @@ end
 
 # Reinterpret an operator-style `Broadcasted` under `NamedTensorStyle`, the broadcast
 # over the operators' states, so the shared `NamedTensorStyle` implementation runs (its
-# `broadcasted_unnamed` already peels each operator operand to its `state` via
-# `unnamed`).
+# `unnamed_linear`/`unnamed_linear_aligned` peel each operator operand to its `state` via `unnamed`).
 function statebroadcasted(bc::Broadcasted{<:NamedTensorOperatorStyle})
     return Broadcasted{NamedTensorStyle{Any}}(bc.f, bc.args, bc.axes)
 end

test/Project.toml

@@ -24,6 +24,10 @@ UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 VectorInterface = "409d34a3-91d5-4945-b6ec-7529ddf182d8"
 WrappedUnions = "325db55a-9c6c-5b90-b1a2-ec87e7a38c44"
 
+[sources.GradedArrays]
+rev = "main"
+url = "https://github.com/ITensor/GradedArrays.jl"
+
 [sources.ITensorBase]
 path = ".."
 
@@ -32,8 +36,8 @@ AbstractTrees = "0.4.5"
 Adapt = "4"
 Aqua = "0.8.9"
 Combinatorics = "1"
-GradedArrays = "0.15"
-ITensorBase = "0.13"
+GradedArrays = "0.16"
+ITensorBase = "0.14"
 ITensorPkgSkeleton = "0.3.42"
 JLArrays = "0.2, 0.3"
 LinearAlgebra = "1.10"
@@ -44,7 +48,7 @@ Random = "1.10"
 SafeTestsets = "0.1"
 StableRNGs = "1"
 Suppressor = "0.2"
-TensorAlgebra = "0.18"
+TensorAlgebra = "0.20"
 TensorKit = "0.17"
 TensorKitSectors = "0.3.9"
 TermInterface = "2"

test/test_gradedarraysext.jl

@@ -1,5 +1,5 @@
 using GradedArrays: U1, sectors
-using ITensorBase: ITensorBase, Index, inds, prime, space
+using ITensorBase: ITensorBase, Index, aligndims, inds, prime, space, unnamed
 using StableRNGs: StableRNG
 using TensorAlgebra: TensorAlgebra, isdual, project, project_aux, tryproject,
     tryproject_aux, unchecked_project, unchecked_project_aux
@@ -76,6 +76,34 @@ using Test: @test, @test_throws, @testset
     @test length(inds(fill(elt(2), U1(1), (), (j,)))) == 2
 end
 
+# Broadcasting over graded (GradedArrays.jl) indices routes the named expression through the
+# `GradedArray` / matricized `FusedGradedMatrix` backend. Linear combinations add block-wise; a sum
+# flattens all-codomain, so a within-split reorder is compared at a common split via `aligndims`.
+@testset "GradedArraysExt broadcasting (eltype = $elt)" for elt in (Float64, ComplexF64)
+    rng = StableRNG(1234)
+    i = Index([U1(0) => 2, U1(1) => 3]; tags = "i")
+    j = Index([U1(0) => 1, U1(1) => 2]; tags = "j")
+    k = Index([U1(-1) => 1, U1(0) => 2]; tags = "k")
+
+    # Flat form (all-codomain, `GradedArray`-backed).
+    a = randn(rng, elt, i, j)
+    b = randn(rng, elt, i, j)
+    @test unnamed(a .+ b) ≈ unnamed(a) + unnamed(b)
+    @test unnamed(2 .* a) ≈ 2 * unnamed(a)
+    @test unnamed(a .- 3 .* b) ≈ unnamed(a) - 3 * unnamed(b)
+
+    # Map form (codomain/domain split, matricized `FusedGradedMatrix` storage).
+    m = randn(rng, elt, (i,), (j,))
+    n = randn(rng, elt, (i,), (j,))
+    @test unnamed(m .+ n) ≈ unnamed(m) + unnamed(n)
+
+    # Within-split reorder still adds correctly (the sum is all-codomain, compared via `aligndims`).
+    mr1 = randn(rng, elt, (i, j), (k,))
+    mr2 = randn(rng, elt, (j, i), (k,))
+    @test unnamed(aligndims(mr1 .+ mr2, (i, j), (k,))) ≈
+        unnamed(mr1) + unnamed(aligndims(mr2, (i, j), (k,)))
+end
+
 # `project_aux` and its siblings derive a named auxiliary leg carrying the operator's flux, so a
 # charge-shifting operator stays symmetry-allowed instead of being projected away. Strict `project`
 # instead projects into exactly the given indices and rejects a surplus axis.

test/test_tensoralgebra.jl

@@ -49,7 +49,7 @@ using Test: @test, @test_broken, @testset
         @test Array(matricize(na, [k, i], [j, l])) == Array(na_pos)
     end
     @testset "unmatricize" begin
-        a, b = namedoneto.((6, 20), ("a", "b"))
+        a, b = namedoneto.((8, 15), ("a", "b"))
         i, j, k, l = namedoneto.((2, 3, 4, 5), ("i", "j", "k", "l"))
         na = randn(elt, a, b)
         # Split all dimensions.

test/test_tensorkitext.jl

@@ -66,13 +66,26 @@ using Test: @test, @test_throws, @testset
         @test TK.space(ref) == TK.space(gc)
         @test ref ≈ gc
 
-        # Linear-combination broadcast lowers to `bipermutedimsopadd!`; element-wise errors.
+        # Linear-combination broadcast lowers to `bipermutedimsopadd!`; a non-linear element-wise
+        # `f.(a)` on a graded tensor errors (graded broadcasting is linear-only).
         b2 = randn(rng, elt, i, j)
         @test unnamed(a + b2) ≈ unnamed(a) + unnamed(b2)
         @test unnamed(2 * a) ≈ 2 * unnamed(a)
         @test unnamed(a .- 3 .* b2) ≈ unnamed(a) - 3 * unnamed(b2)
         @test_throws ErrorException sin.(a)
 
+        # Named broadcasting aligns operands by name within their codomain/domain split, so a within-split
+        # reorder of a multi-leg operand still adds correctly (compared at a common split via `aligndims`).
+        mr1 = randn(rng, elt, (i, j), (k,))
+        mr2 = randn(rng, elt, (j, i), (k,))
+        @test unnamed(aligndims(mr1 .+ mr2, (i, j), (k,))) ≈
+            unnamed(mr1) + unnamed(aligndims(mr2, (i, j), (k,)))
+        # Adding across an incompatible split (a shared leg in the codomain of one operand and the domain
+        # of the other) has mismatched axes and errors.
+        cs1 = randn(rng, elt, (i,), (j,))
+        cs2 = randn(rng, elt, (j,), (i,))
+        @test_throws DimensionMismatch cs1 .+ cs2
+
         # Factorizations reconstruct the tensor (lowered through matricize / MatrixAlgebraKit).
         # Checked by full contraction to a scalar, which is bipartition-independent.
         a3 = randn(rng, elt, i, j, k)