Getting lazy to work

Author: JoeyT1994

Project.toml

@@ -1,9 +1,10 @@
 name = "ITensorNetworksNext"
 uuid = "302f2e75-49f0-4526-aef7-d8ba550cb06c"
 authors = ["ITensor developers <[email protected]> and contributors"]
-version = "0.1.7"
+version = "0.1.9"
 
 [deps]
+AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 BackendSelection = "680c2d7c-f67a-4cc9-ae9c-da132b1447a5"
 DataGraphs = "b5a273c3-7e6c-41f6-98bd-8d7f1525a36a"
@@ -23,9 +24,11 @@ TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2"
 [extensions]
 ITensorNetworksNextTensorOperationsExt = "TensorOperations"
 TermInterface = "8ea1fca8-c5ef-4a55-8b96-4e9afe9c9a3c"
+TypeParameterAccessors = "7e5a90cf-f82e-492e-a09b-e3e26432c138"
 WrappedUnions = "325db55a-9c6c-5b90-b1a2-ec87e7a38c44"
 
 [compat]
+AbstractTrees = "0.4.5"
 Adapt = "4.3"
 BackendSelection = "0.1.6"
 DataGraphs = "0.2.7"
@@ -39,6 +42,7 @@ NamedGraphs = "0.6.9, 0.7"
 SimpleTraits = "0.9.5"
 SplitApplyCombine = "1.2.3"
 TermInterface = "2"
+TypeParameterAccessors = "0.4.4"
 WrappedUnions = "0.3"
 TensorOperations = "5.3.1"
 julia = "1.10"

ext/ITensorNetworksNextTensorOperationsExt/ITensorNetworksNextTensorOperationsExt.jl

@@ -1,16 +1,17 @@
 module ITensorNetworksNextTensorOperationsExt
 
 using BackendSelection: @Algorithm_str, Algorithm
-using ITensorBase: inds, dim
+using ITensorBase: inds
 using ITensorNetworksNext: ITensorNetworksNext
+using ITensorNetworksNext.LazyNamedDimsArrays: nested_array_to_lazy_multiply
 using TensorOperations: TensorOperations, optimaltree
 
 function ITensorNetworksNext.contraction_sequence(::Algorithm"optimal", tn::Vector{<:AbstractArray})
     network = collect.(inds.(tn))
     #Converting dims to Float64 to minimize overflow issues
-    inds_to_dims = Dict(i => Float64(dim(i)) for i in unique(reduce(vcat, network)))
+    inds_to_dims = Dict(i => Float64(length(i)) for i in unique(reduce(vcat, network)))
     seq, _ = optimaltree(network, inds_to_dims)
-    return seq
+    return nested_array_to_lazy_multiply(seq)
 end
 
 end

src/contractnetwork.jl

@@ -1,4 +1,5 @@
 using BackendSelection: @Algorithm_str, Algorithm
+using ITensorNetworksNext.LazyNamedDimsArrays: LazyNamedDimsArray, nested_array_to_lazy_multiply, substitute_lazy, materialize
 
 default_contract_alg = nothing
 
@@ -10,26 +11,20 @@ function set_default_kwargs(alg::Algorithm"exact")
 end
 
 function contraction_sequence(::Algorithm"leftassociative", tn::Vector{<:AbstractArray})
-    return Any[i for i in 1:length(tn)]
+    return nested_array_to_lazy_multiply(collect.(1:length(tn)))
 end
 
 function contraction_sequence(tn::Vector{<:AbstractArray}; alg = default_sequence_alg)
     return contraction_sequence(Algorithm(alg), tn)
 end
 
-# Internal recursive worker
-function recursive_contractnetwork(tn::Union{AbstractVector, AbstractNamedDimsArray})
-    tn isa AbstractVector && return prod(recursive_contractnetwork, tn)
-    return tn
-end
-
-# Recursive worker for ordering the tensors according to the sequence
-rearrange(tn::Vector{<:AbstractArray}, i::Integer) = tn[i]
-rearrange(tn::Vector{<:AbstractArray}, v::AbstractVector) = [rearrange(tn, s) for s in v]
-
 function contractnetwork(alg::Algorithm"exact", tn::Vector{<:AbstractArray})
     contract_sequence = isa(alg.sequence, String) ? contraction_sequence(tn; alg = alg.sequence) : sequence
-    return recursive_contractnetwork(rearrange(tn, contract_sequence))
+    @show contract_sequence
+    @show materialize(contract_sequence)
+    contract_sequence = substitute_lazy(contract_sequence, Dict(i => lazy(tn[i]) for i in 1:length(tn)))
+    @show contract_sequence
+    #return materialize(substitute_lazy(contract_sequence, Dict(i => tn[i] for i in 1:length(tn))))
 end
 
 function contractnetwork(alg::Algorithm"exact", tn::AbstractTensorNetwork)