From 119b408442d5f45ac15bfa138227bda18f984902 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Tue, 28 Jul 2026 20:54:17 -0400 Subject: [PATCH] Promote lazy tensors and re-partition operators Extends the operator promotion from #235 to lazy tensors: a LazyNamedTensor promotes with lazy or eager tensors by promoting the leaf type, so a lazy plain tensor and an operator promote to a lazy operator. Also makes operator() on an existing operator reassign the pairing rather than nesting. Together these let contract_network homogenize networks mixing lazy/eager, plain/operator operands. --- Project.toml | 2 +- src/lazyitensors/lazyitensor.jl | 28 ++++++++++++++++++++++++++ src/namedtensoroperator.jl | 5 +++++ test/test_lazyitensors.jl | 35 ++++++++++++++++++++++++++++++--- test/test_operator.jl | 11 +++++++++++ 5 files changed, 77 insertions(+), 4 deletions(-) diff --git a/Project.toml b/Project.toml index e874e497..34bf4aee 100644 --- a/Project.toml +++ b/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 and contributors"] [workspace] diff --git a/src/lazyitensors/lazyitensor.jl b/src/lazyitensors/lazyitensor.jl index 6046b6aa..975e11f4 100644 --- a/src/lazyitensors/lazyitensor.jl +++ b/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 diff --git a/src/namedtensoroperator.jl b/src/namedtensoroperator.jl index f97de939..92312aca 100644 --- a/src/namedtensoroperator.jl +++ b/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` diff --git a/test/test_lazyitensors.jl b/test/test_lazyitensors.jl index 4f560d8c..06412c5e 100644 --- a/test/test_lazyitensors.jl +++ b/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 diff --git a/test/test_operator.jl b/test/test_operator.jl index 87f24e8f..cb377d5d 100644 --- a/test/test_operator.jl +++ b/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