Support broadcasting on operators and dot on graded named arrays

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 <support@itensor.org> and contributors"]
 
 [workspace]

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

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)

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)