[WIP] Make conj of a named tensor lazy

Project.toml

@@ -1,6 +1,6 @@
 name = "ITensorBase"
 uuid = "4795dd04-0d67-49bb-8f44-b89c448a1dc7"
-version = "0.16.4"
+version = "0.16.5"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
 
 [workspace]
@@ -9,6 +9,7 @@ projects = ["benchmark", "dev", "docs", "examples", "test"]
 [deps]
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
+GradedArrays = "bc96ca6e-b7c8-4bb6-888e-c93f838762c2"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MatrixAlgebraKit = "6c742aac-3347-4629-af66-fc926824e5e4"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -19,11 +20,18 @@ VectorInterface = "409d34a3-91d5-4945-b6ec-7529ddf182d8"
 
 [weakdeps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
-GradedArrays = "bc96ca6e-b7c8-4bb6-888e-c93f838762c2"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 TensorKit = "07d1fe3e-3e46-537d-9eac-e9e13d0d4cec"
 TensorKitSectors = "13a9c161-d5da-41f0-bcbd-e1a08ae0647f"
 
+[sources.GradedArrays]
+rev = "mf/conjarray-contract"
+url = "https://github.com/ITensor/GradedArrays.jl"
+
+[sources.TensorAlgebra]
+rev = "mf/lazy-conjarray"
+url = "https://github.com/ITensor/TensorAlgebra.jl"
+
 [extensions]
 ITensorBaseAdaptExt = "Adapt"
 ITensorBaseGradedArraysExt = ["GradedArrays", "TensorKitSectors"]
@@ -34,13 +42,13 @@ ITensorBaseTensorKitExt = "TensorKit"
 Accessors = "0.1.39"
 Adapt = "4.1.1"
 ArrayLayouts = "1.11"
-GradedArrays = "0.17.3"
+GradedArrays = "0.17.8"
 LinearAlgebra = "1.10"
 MatrixAlgebraKit = "0.2, 0.3, 0.4, 0.5, 0.6"
 Mooncake = "0.4.202, 0.5"
 Random = "1.10"
 SimpleTraits = "0.9.4"
-TensorAlgebra = "0.23.2"
+TensorAlgebra = "0.23.4"
 TensorKit = "0.17"
 TensorKitSectors = "0.3.9"
 UUIDs = "1.10"

src/abstractnamedtensor.jl

@@ -1,6 +1,6 @@
 using LinearAlgebra: LinearAlgebra
 using Random: Random
-using TensorAlgebra: TensorAlgebra, permuteddims, zero!
+using TensorAlgebra: TensorAlgebra, conjed, permuteddims, zero!
 
 # Some of the interface is inspired by:
 # https://github.com/NamedTensor/ITensors.jl
@@ -183,14 +183,16 @@ Base.zero(a::AbstractNamedTensor) = namedtensorof(a, zero(unnamed(a)))
 # `AbstractArray` fallback did through `axes`).
 Base.CartesianIndices(a::AbstractNamedTensor) = CartesianIndices(axes(a))
 
-# Eager conjugation reuses the linear-combination broadcast machinery: `conj.(a)` lowers to a
-# `ConjBroadcasted` leaf that the op primitives absorb (folding `conj` into their `op` flag) and
-# materializes into an all-codomain destination with dualized axes. This keeps one conjugation
-# mechanism shared by the lazy `conj.` path and eager `conj`, so a non-`AbstractArray` backend
-# (a `TensorMap`) needs no separate eager hook. The default `AbstractArray` `conj` would instead
-# map over elements without dualizing the axes, silently changing the contraction convention for
-# graded tensors.
-Base.conj(a::AbstractNamedTensor) = conj.(a)
+# Conjugation is lazy for an `AbstractArray` backing: `conjed` wraps it in a `ConjArray`, which
+# contraction, permutation and matricization absorb by folding `conj` into their `op`, so the
+# conjugate is never materialized on its way into them. A non-`AbstractArray` backing (a
+# `TensorMap`) conjugates eagerly through the `conj.(a)` broadcast instead, which materializes into
+# an all-codomain destination with dualized axes. The default `AbstractArray` `conj` would map over
+# elements without dualizing the axes, silently changing the contraction convention for graded
+# tensors.
+Base.conj(a::AbstractNamedTensor) = conj_namedtensor(a, unnamed(a))
+conj_namedtensor(a, unnamed_a::AbstractArray) = namedtensorof(a, conjed(unnamed_a))
+conj_namedtensor(a, unnamed_a) = conj.(a)
 
 # `LinearAlgebra.normalize` infers result eltype via `typeof(first(a)/nrm)`, which
 # scalar-indexes block-structured storage. `a / norm(a, p)` already preserves names.

test/test_gradedarraysext.jl

@@ -1,10 +1,11 @@
 using GradedArrays: U1, fU1, sectors
 using ITensorBase:
     ITensorBase, ITensor, Index, align, inds, name, prime, space, unnamed, uuid
+using LinearAlgebra: LinearAlgebra
 using StableRNGs: StableRNG
-using TensorAlgebra: TensorAlgebra, dual, isdual, matricize, project, project_aux,
-    tryproject, tryproject_aux, twist!, unchecked_project, unchecked_project_aux,
-    unmatricize
+using TensorAlgebra: TensorAlgebra, ConjArray, dual, isdual, matricize, project,
+    project_aux, tryproject, tryproject_aux, twist!, unchecked_project,
+    unchecked_project_aux, unmatricize
 using TensorKitSectors: FermionNumber
 using Test: @test, @test_throws, @testset
 
@@ -104,6 +105,25 @@ end
         unnamed(mr1) + unnamed(align(mr2, (i, j), (k,)))
 end
 
+# `conj` of a graded named tensor is lazy and folds into the contraction, so it matches the eager
+# conjugate, fermion signs included.
+@testset "GradedArraysExt lazy conj (eltype = $elt)" for elt in (Float64, ComplexF64)
+    rng = StableRNG(1234)
+    i = Index([fU1(0) => 2, fU1(1) => 3]; tags = "i")
+    j = Index([fU1(0) => 1, fU1(1) => 2]; tags = "j")
+    k = Index([fU1(-1) => 1, fU1(0) => 2]; tags = "k")
+    a = randn(rng, elt, i, j, k)
+    c = randn(rng, elt, i, j)
+    @test unnamed(conj(a)) isa ConjArray
+    @test conj(a) * c ≈ conj.(a) * c
+    @test c * conj(a) ≈ c * conj.(a)
+    @test conj(a) * a ≈ conj.(a) * a
+    # The other operand shares the dualized axes of `conj(a)`, with its names in a different order.
+    d = conj.(randn(rng, elt, k, i, j))
+    @test LinearAlgebra.dot(conj(a), d) ≈ LinearAlgebra.dot(conj.(a), d)
+    @test LinearAlgebra.dot(d, conj(a)) ≈ LinearAlgebra.dot(d, conj.(a))
+end
+
 # `matricize` fuses a tensor's codomain/domain split into an unnamed matrix and `unmatricize`
 # splits one back out over named indices. The domain index is stored dualized while it is given
 # to `unmatricize` codomain-facing, so a graded backend is where that convention is visible.

test/test_namedtensor_basics.jl

@@ -5,7 +5,7 @@ using ITensorBase: @names, AbstractNamedTensor, Name, NameMismatch, Named,
     product, rename, setnames, unname, unnamed, unnamedtype
 using LinearAlgebra: LinearAlgebra
 using Random: default_rng
-using TensorAlgebra: datatype
+using TensorAlgebra: ConjArray, datatype
 using Test: @test, @test_throws, @testset
 using VectorInterface: scalartype
 
@@ -386,6 +386,21 @@ end
         @test unnamed(ca) == conj(unnamed(a))
         @test names(ca) == names(a)
 
+        # `conj` is lazy, and contracting the lazy conjugate matches the eager one.
+        @test unnamed(ca) isa ConjArray
+        @test unnamed(conj(ca)) === unnamed(a)
+        b = randn(elt, j, NamedOneTo(4, "k"))
+        @test ca * b == conj.(a) * b
+        # Named `dot` goes through `dotperm`, so the names may come in a different order.
+        y = randn(elt, j, i)
+        @test LinearAlgebra.dot(ca, y) ≈ LinearAlgebra.dot(conj.(a), y)
+        @test LinearAlgebra.dot(y, ca) ≈ LinearAlgebra.dot(y, conj.(a))
+
+        # The lazy conjugate is a view, so writing into it writes the conjugate into `a`.
+        x = randn(elt, i, j)
+        @test fill!(conj(x), zero(elt)) isa AbstractNamedTensor
+        @test iszero(unnamed(x))
+
         # `fill!` forwards to the underlying storage.
         b = randn(elt, i, j)
         @test fill!(b, zero(elt)) === b