diff --git a/Project.toml b/Project.toml index b4f3ec1..fdded47 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "ITensorBase" uuid = "4795dd04-0d67-49bb-8f44-b89c448a1dc7" -version = "0.13.13" +version = "0.13.14" authors = ["ITensor developers and contributors"] [workspace] @@ -45,14 +45,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.19.3" TensorKit = "0.17" TensorKitSectors = "0.3.9" TermInterface = "2" diff --git a/docs/Project.toml b/docs/Project.toml index 03074f5..22c7796 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -16,5 +16,5 @@ ITensorBase = "0.13" ITensorFormatter = "0.2.27" Literate = "2" MatrixAlgebraKit = "0.2, 0.3, 0.4, 0.5, 0.6" -TensorAlgebra = "0.18" +TensorAlgebra = "0.19.3" Test = "1.10" diff --git a/src/broadcast.jl b/src/broadcast.jl index 5f97907..2e11211 100644 --- a/src/broadcast.jl +++ b/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 diff --git a/test/Project.toml b/test/Project.toml index 4f0da76..f195685 100644 --- a/test/Project.toml +++ b/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,7 +36,7 @@ AbstractTrees = "0.4.5" Adapt = "4" Aqua = "0.8.9" Combinatorics = "1" -GradedArrays = "0.15" +GradedArrays = "0.16" ITensorBase = "0.13" ITensorPkgSkeleton = "0.3.42" JLArrays = "0.2, 0.3" @@ -44,7 +48,7 @@ Random = "1.10" SafeTestsets = "0.1" StableRNGs = "1" Suppressor = "0.2" -TensorAlgebra = "0.18" +TensorAlgebra = "0.19.3" TensorKit = "0.17" TensorKitSectors = "0.3.9" TermInterface = "2" diff --git a/test/test_gradedarraysext.jl b/test/test_gradedarraysext.jl index 6e7095e..43867cf 100644 --- a/test/test_gradedarraysext.jl +++ b/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. diff --git a/test/test_tensorkitext.jl b/test/test_tensorkitext.jl index ebe6258..c439454 100644 --- a/test/test_tensorkitext.jl +++ b/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)