Promote lazy tensors and re-partition operators

Project.toml

@@ -1,6 +1,6 @@
 name = "ITensorBase"
 uuid = "4795dd04-0d67-49bb-8f44-b89c448a1dc7"
-version = "0.13.11"
+version = "0.13.12"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
 
 [workspace]

src/lazyitensors/lazyitensor.jl

@@ -22,6 +22,34 @@ lazy(a::LazyNamedTensor) = a
 lazy(a::AbstractNamedTensor) = LazyNamedTensor(a)
 lazy(a::Mul{<:LazyNamedTensor}) = LazyNamedTensor(a)
 
+# Promotion of lazy tensors, by promoting the wrapped (leaf) parent type. A lazy tensor promotes
+# with another lazy tensor, or with an eager tensor, to the lazy tensor whose leaf type is the
+# promotion of the leaves. So a lazy plain tensor and an operator promote to a lazy operator,
+# reusing the eager `NamedTensor` / `NamedTensorOperator` promotion at the leaves. This lets
+# `contract_network` homogenize a network that mixes lazy and eager, plain and operator operands
+# (a norm-network vertex is a lazy `ket * conj(bra)` product) without materializing anything.
+function Base.promote_rule(
+        ::Type{LazyNamedTensor{D, A1}}, ::Type{LazyNamedTensor{D, A2}}
+    ) where {D, A1, A2}
+    return LazyNamedTensor{D, promote_type(A1, A2)}
+end
+function Base.promote_rule(
+        ::Type{LazyNamedTensor{D, A}}, ::Type{T}
+    ) where {D, A, T <: AbstractNamedTensor{D}}
+    return LazyNamedTensor{D, promote_type(A, T)}
+end
+# `convert` mirrors the promotion: wrap an eager tensor as lazy (converting it to the target leaf
+# type first), or rebuild a lazy product with each leaf converted to the target leaf type.
+function Base.convert(::Type{LazyNamedTensor{D, A}}, a::AbstractNamedTensor{D}) where {D, A}
+    return lazy(convert(A, a))
+end
+function Base.convert(
+        ::Type{LazyNamedTensor{D, A2}}, a::LazyNamedTensor{D, A1}
+    ) where {D, A1, A2}
+    iscall(a) || return lazy(convert(A2, unwrap(a)))
+    return lazy(Mul(map(arg -> convert(LazyNamedTensor{D, A2}, arg), arguments(a))))
+end
+
 dimnames(a::LazyNamedTensor) = dimnames_lazy(a)
 inds(a::LazyNamedTensor) = inds_lazy(a)
 # `axes` is computed from `inds_lazy` rather than the generic `unnamed`-based fallback

src/namedtensoroperator.jl

@@ -389,6 +389,11 @@ end
 function operator(a::AbstractNamedTensor, output, input)
     return NamedTensorOperator(a, name.(output), name.(input))
 end
+# An operator's output/input pairing is a view over its state, so `operator` on an existing operator
+# reassigns the pairing rather than nesting another operator around it.
+function operator(a::NamedTensorOperator, output, input)
+    return operator(state(a), output, input)
+end
 
 # A plain tensor is a trivial (empty-pairing) operator, so an operator is the promotion of a
 # non-operator tensor. `convert` wraps a plain tensor as a trivial operator and `promote_rule`

test/test_lazyitensors.jl

@@ -1,8 +1,8 @@
 using AbstractTrees: AbstractTrees, print_tree, printnode
 using Base.Broadcast: materialize
-using ITensorBase: @names, Greedy, LazyNamedTensor, Mul, NamedTensor, SymbolicNamedTensor,
-    dimnames, inds, ismul, lazy, nameddims, namedoneto, optimize_evaluation_order,
-    substitute, symnameddims
+using ITensorBase: @names, Greedy, LazyNamedTensor, Mul, NamedTensor, NamedTensorOperator,
+    SymbolicNamedTensor, dimnames, inds, inputnames, ismul, lazy, nameddims, namedoneto,
+    operator, optimize_evaluation_order, outputnames, state, substitute, symnameddims
 using OMEinsumContractionOrders: ExhaustiveSearch, GreedyMethod, TreeSA
 using TermInterface: arguments, arity, children, head, iscall, isexpr, maketerm, operation,
     sorted_arguments, sorted_children
@@ -136,3 +136,32 @@ using WrappedUnions: unwrap
         @test issetequal(dimnames(ordered), dimnames(flat))
     end
 end
+
+@testset "lazy operator promotion" begin
+    i, j = namedoneto.(2, (:i, :j))
+    p = randn(i, j)                          # eager plain
+    o = operator(randn(i, j), (i,), (j,))    # eager operator
+    lp = lazy(p)                             # lazy plain
+    lo = lazy(o)                             # lazy operator
+    P, O, LP, LO = typeof(p), typeof(o), typeof(lp), typeof(lo)
+
+    # A lazy tensor promotes with a lazy or an eager tensor by promoting the leaf type, so any mix
+    # containing a lazy operand or an operator climbs to the lazy operator `LO`, while an all-plain
+    # network stays eager.
+    @test promote_type(P, O) == O
+    @test promote_type(P, LP) == LP
+    @test promote_type(P, LO) == LO
+    @test promote_type(O, LP) == LO
+    @test promote_type(O, LO) == LO
+    @test promote_type(LP, LO) == LO
+    @test promote_type(P, P) == P
+
+    # `convert` wraps an eager tensor as lazy (converting the leaf), and rebuilds a lazy product with
+    # each leaf converted to the target leaf type.
+    @test convert(LP, p) isa LP
+    @test convert(LO, o) isa LO
+    clo = convert(LO, lp)
+    @test clo isa LO
+    @test materialize(clo) isa NamedTensorOperator
+    @test isempty(outputnames(materialize(clo))) && isempty(inputnames(materialize(clo)))
+end

test/test_operator.jl

@@ -498,3 +498,14 @@ end
     @test eltype(v) == O
     @test all(x -> x isa NamedTensorOperator, v)
 end
+
+@testset "operator re-partition" begin
+    o = operator(randn(2, 2), ("i",), ("j",))
+    # `operator` on an existing operator reassigns the pairing over the same state rather than
+    # nesting another operator around it.
+    o2 = operator(o, ("j",), ("i",))
+    @test o2 isa NamedTensorOperator
+    @test state(o2) === state(o)
+    @test outputnames(o2) == ["j"]
+    @test inputnames(o2) == ["i"]
+end