From de6890b521fe7437f80611aa281008ae20e216bd Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Thu, 18 Jun 2026 12:29:08 -0400 Subject: [PATCH] Support broadcasting on operators and dot on graded named arrays Lets operator-backed and graded-backed named arrays go through linear algebra without falling into the generic `AbstractArray` paths, which assume a statically known `ndims` or use scalar indexing. `Base.broadcastable` on `AbstractNamedDimsOperator` peels an operator to its `state`, so broadcasting and the array operations that lower to it (`+`, `-`, scalar multiplication) run on the underlying state. Without it they threw `TypeError: expected Int64, got Type{Any}` when the operator's parent is dynamically ranked, as it is when backed by an `ITensor`. The result is a bare state, dropping the operator wrapper, matching `*`. `dot` becomes the full contraction `(conj(a1) * a2)[]`, which aligns the legs by name and reduces through the contraction backend. The generic definition scalar-indexes, which graded arrays reject. Together these let ITensorNetworksNext run belief propagation with operator-valued messages, where message normalization and the convergence check exercise both paths. --- Project.toml | 2 +- src/linearalgebra.jl | 5 +---- src/nameddimsoperator.jl | 10 ++++++++++ test/test_operator.jl | 25 +++++++++++++++++++++++++ 4 files changed, 37 insertions(+), 5 deletions(-) diff --git a/Project.toml b/Project.toml index d8ae3852..3a3cf01a 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "NamedDimsArrays" uuid = "60cbd0c0-df58-4cb7-918c-6f5607b73fde" -version = "0.15.9" +version = "0.15.10" authors = ["ITensor developers and contributors"] [workspace] diff --git a/src/linearalgebra.jl b/src/linearalgebra.jl index 1da08b34..0921e1d9 100644 --- a/src/linearalgebra.jl +++ b/src/linearalgebra.jl @@ -34,8 +34,5 @@ end # https://github.com/JuliaLang/LinearAlgebra.jl/blob/3a4fdad7f608928ecb4b41e76b1e9ecacd058444/src/generic.jl#L919-L1009 # which isn't friendly for NamedDimsArrays wrapping GPU arrays. function LA.dot(a1::AbstractNamedDimsArray, a2::AbstractNamedDimsArray) - # TODO: Should we define: - # `TensorAlgebra.permdot(denamed(a1), denamed(a2), genperm(dimnames(a1), dimnames(a2)))` - # in TensorAlgebra.jl and use that here? - return LA.dot(denamed(a1), denamed(a2, dimnames(a1))) + return (conj(a1) * a2)[] end diff --git a/src/nameddimsoperator.jl b/src/nameddimsoperator.jl index 2f6e4fca..2047f6b5 100644 --- a/src/nameddimsoperator.jl +++ b/src/nameddimsoperator.jl @@ -117,6 +117,16 @@ Base.:*(a::AbstractNamedDimsOperator, b::AbstractNamedDimsOperator) = state(a) * Base.:*(a::AbstractNamedDimsOperator, b::AbstractNamedDimsArray) = state(a) * state(b) Base.:*(a::AbstractNamedDimsArray, b::AbstractNamedDimsOperator) = state(a) * state(b) +# Broadcasting peels an operator to its `state` before the broadcast style is +# computed. This covers the array operations that lower to broadcasting too (`+`, +# `-`, scalar multiplication), and avoids the generic style path reading `ndims` +# off the type, which throws for a dynamically-ranked parent such as an `ITensor`. +# As a shortcut the result is a bare state, dropping the operator wrapper, matching +# `*`; keeping it an operator (carrying the codomain/domain bijection) is future +# work that first needs the mixed cases settled (`state + operator`, and +# `operator + operator` where names match but the codomain/domain split does not). +Base.broadcastable(a::AbstractNamedDimsOperator) = state(a) + for f in MATRIX_FUNCTIONS @eval begin function Base.$f(a::AbstractNamedDimsOperator) diff --git a/test/test_operator.jl b/test/test_operator.jl index 6c1e9e52..f162e033 100644 --- a/test/test_operator.jl +++ b/test/test_operator.jl @@ -109,6 +109,31 @@ end @test all(0 .≤ denamed(state(op)) .≤ 1) end +@testset "operator linear algebra peels to state" begin + # `+`, `-`, and scalar multiplication lower to broadcasting, which peels an + # operator to its underlying state (via `broadcastable`), so the result drops + # the operator wrapper, matching `*`. Without that, the generic broadcast path + # reads `ndims` off the type and throws for a dynamically-ranked parent such as + # an `ITensor`. A future version may keep the result an operator instead; this + # pins the current contract. + o = operator(randn(2, 2), ("i'",), ("i",)) + s = state(o) + nms = ("i'", "i") + + for r in (o + o, o - o, -o, 2 * o, o * 2, 2 .* o, o .* 2) + @test r isa NamedDimsArray + @test !(r isa NamedDimsOperator) + end + + @test dename(o + o, nms) ≈ 2 .* dename(s, nms) + @test all(iszero, dename(o - o, nms)) + @test dename(-o, nms) ≈ -dename(s, nms) + @test dename(2 * o, nms) ≈ 2 .* dename(s, nms) + @test dename(o * 2, nms) ≈ 2 .* dename(s, nms) + @test dename(2 .* o, nms) ≈ 2 .* dename(s, nms) + @test dename(o .* 2, nms) ≈ 2 .* dename(s, nms) +end + @testset "gram_eigh_full on AbstractNamedDimsOperator" begin n = 5 B = randn(n, n)