Contraction sequence finding (#218)

Author: JoeyT1994

Project.toml

@@ -1,7 +1,7 @@
 name = "ITensorNetworks"
 uuid = "2919e153-833c-4bdc-8836-1ea460a35fc7"
 authors = ["Matthew Fishman <[email protected]>, Joseph Tindall <[email protected]> and contributors"]
-version = "0.11.27"
+version = "0.12.0"
 
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
@@ -41,13 +41,15 @@ EinExprs = "b1794770-133b-4de1-afb4-526377e9f4c5"
 GraphsFlows = "06909019-6f44-4949-96fc-b9d9aaa02889"
 OMEinsumContractionOrders = "6f22d1fd-8eed-4bb7-9776-e7d684900715"
 Observers = "338f10d5-c7f1-4033-a7d1-f9dec39bcaa0"
+TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2"
 
 [extensions]
 ITensorNetworksAdaptExt = "Adapt"
 ITensorNetworksEinExprsExt = "EinExprs"
 ITensorNetworksGraphsFlowsExt = "GraphsFlows"
 ITensorNetworksOMEinsumContractionOrdersExt = "OMEinsumContractionOrders"
 ITensorNetworksObserversExt = "Observers"
+ITensorNetworksTensorOperationsExt = "TensorOperations"
 
 [compat]
 AbstractTrees = "0.4.4"
@@ -80,6 +82,7 @@ SplitApplyCombine = "1.2"
 StaticArrays = "1.5.12"
 StructWalk = "0.2"
 Suppressor = "0.2"
+TensorOperations = "5.1.4"
 TimerOutputs = "0.5.22"
 TupleTools = "1.4"
 julia = "1.10"
@@ -90,6 +93,7 @@ EinExprs = "b1794770-133b-4de1-afb4-526377e9f4c5"
 GraphsFlows = "06909019-6f44-4949-96fc-b9d9aaa02889"
 OMEinsumContractionOrders = "6f22d1fd-8eed-4bb7-9776-e7d684900715"
 Observers = "338f10d5-c7f1-4033-a7d1-f9dec39bcaa0"
+TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]

ext/ITensorNetworksEinExprsExt/ITensorNetworksEinExprsExt.jl

@@ -2,10 +2,15 @@ module ITensorNetworksEinExprsExt
 
 using ITensors: Index, ITensor, @Algorithm_str, inds, noncommoninds
 using ITensorNetworks:
-  ITensorNetworks, ITensorNetwork, vertextype, vertex_data, contraction_sequence
+  ITensorNetworks,
+  ITensorList,
+  ITensorNetwork,
+  vertextype,
+  vertex_data,
+  contraction_sequence
 using EinExprs: EinExprs, EinExpr, einexpr, SizedEinExpr
 
-function to_einexpr(ts::Vector{ITensor})
+function to_einexpr(ts::ITensorList)
   IndexType = Any
 
   tensor_exprs = EinExpr{IndexType}[]
@@ -21,7 +26,7 @@ function to_einexpr(ts::Vector{ITensor})
   return SizedEinExpr(sum(tensor_exprs; skip=externalinds_tn), inds_dims)
 end
 
-function tensor_inds_to_vertex(ts::Vector{ITensor})
+function tensor_inds_to_vertex(ts::ITensorList)
   IndexType = Any
   VertexType = Int
 
@@ -36,7 +41,7 @@ function tensor_inds_to_vertex(ts::Vector{ITensor})
 end
 
 function ITensorNetworks.contraction_sequence(
-  ::Algorithm"einexpr", tn::Vector{ITensor}; optimizer=EinExprs.Exhaustive()
+  ::Algorithm"einexpr", tn::ITensorList; optimizer=EinExprs.Exhaustive()
 )
   expr = to_einexpr(tn)
   path = einexpr(optimizer, expr)

ext/ITensorNetworksOMEinsumContractionOrdersExt/ITensorNetworksOMEinsumContractionOrdersExt.jl

@@ -1,6 +1,6 @@
 module ITensorNetworksOMEinsumContractionOrdersExt
 using DocStringExtensions: TYPEDSIGNATURES
-using ITensorNetworks: ITensorNetworks
+using ITensorNetworks: ITensorNetworks, ITensorList
 using ITensors: ITensors, Index, ITensor, inds
 using NDTensors: dim
 using NDTensors.AlgorithmSelection: @Algorithm_str
@@ -9,8 +9,6 @@ using OMEinsumContractionOrders: OMEinsumContractionOrders
 # OMEinsumContractionOrders wrapper for ITensors
 # Slicing is not supported, because it might require extra work to slice an `ITensor` correctly.
 
-const ITensorList = Union{Vector{ITensor},Tuple{Vararg{ITensor}}}
-
 # infer the output tensor labels
 # TODO: Use `symdiff` instead.
 function infer_output(inputs::AbstractVector{<:AbstractVector{<:Index}})
@@ -126,7 +124,9 @@ Optimize the einsum contraction pattern using the simulated annealing on tensor
 ### References
 * [Recursive Multi-Tensor Contraction for XEB Verification of Quantum Circuits](https://arxiv.org/abs/2108.05665)
 """
-function ITensorNetworks.contraction_sequence(::Algorithm"tree_sa", tn; kwargs...)
+function ITensorNetworks.contraction_sequence(
+  ::Algorithm"tree_sa", tn::ITensorList; kwargs...
+)
   return optimize_contraction_sequence(
     tn; optimizer=OMEinsumContractionOrders.TreeSA(; kwargs...)
   )
@@ -153,7 +153,9 @@ Then finds the contraction order inside each group with the greedy search algori
 ### References
 * [Hyper-optimized tensor network contraction](https://arxiv.org/abs/2002.01935)
 """
-function ITensorNetworks.contraction_sequence(::Algorithm"sa_bipartite", tn; kwargs...)
+function ITensorNetworks.contraction_sequence(
+  ::Algorithm"sa_bipartite", tn::ITensorList; kwargs...
+)
   return optimize_contraction_sequence(
     tn; optimizer=OMEinsumContractionOrders.SABipartite(; kwargs...)
   )
@@ -177,7 +179,9 @@ Then finds the contraction order inside each group with the greedy search algori
 * [Hyper-optimized tensor network contraction](https://arxiv.org/abs/2002.01935)
 * [Simulating the Sycamore quantum supremacy circuits](https://arxiv.org/abs/2103.03074)
 """
-function ITensorNetworks.contraction_sequence(::Algorithm"kahypar_bipartite", tn; kwargs...)
+function ITensorNetworks.contraction_sequence(
+  ::Algorithm"kahypar_bipartite", tn::ITensorList; kwargs...
+)
   return optimize_contraction_sequence(
     tn; optimizer=OMEinsumContractionOrders.KaHyParBipartite(; kwargs...)
   )

ext/ITensorNetworksTensorOperationsExt/ITensorNetworksTensorOperationsExt.jl

@@ -0,0 +1,15 @@
+module ITensorNetworksTensorOperationsExt
+
+using ITensors: ITensors, ITensor, dim, inds
+using ITensorNetworks: ITensorNetworks, ITensorList
+using NDTensors.AlgorithmSelection: @Algorithm_str
+using TensorOperations: TensorOperations, optimaltree
+
+function ITensorNetworks.contraction_sequence(::Algorithm"optimal", tn::ITensorList)
+  network = collect.(inds.(tn))
+  inds_to_dims = Dict(i => dim(i) for i in unique(reduce(vcat, network)))
+  seq, _ = optimaltree(network, inds_to_dims)
+  return seq
+end
+
+end

src/contract.jl

@@ -1,5 +1,4 @@
 using ITensors: ITensor, scalar
-using ITensors.ContractionSequenceOptimization: deepmap
 using ITensors.NDTensors: NDTensors, Algorithm, @Algorithm_str, contract
 using LinearAlgebra: normalize!
 using NamedGraphs: NamedGraphs

src/contraction_sequences.jl

@@ -1,22 +1,32 @@
 using Graphs: vertices
-using ITensors: ITensor, contract
-using ITensors.ContractionSequenceOptimization: deepmap, optimal_contraction_sequence
+using ITensors: ITensor
 using ITensors.NDTensors: Algorithm, @Algorithm_str
 using NamedGraphs.Keys: Key
 using NamedGraphs.OrdinalIndexing: th
 
-function contraction_sequence(tn::Vector{ITensor}; alg="optimal", kwargs...)
+const ITensorList = Union{Vector{ITensor},Tuple{Vararg{ITensor}}}
+
+function contraction_sequence(tn::ITensorList; alg="optimal", kwargs...)
   return contraction_sequence(Algorithm(alg), tn; kwargs...)
 end
 
+function contraction_sequence(alg::Algorithm, tn::ITensorList)
+  return throw(
+    ArgumentError(
+      "Algorithm $alg isn't defined for contraction sequence finding. Try loading a backend package like 
+        TensorOperations.jl or OMEinsumContractionOrders.jl.",
+    ),
+  )
+end
+
+function deepmap(f, tree; filter=(x -> x isa AbstractArray))
+  return filter(tree) ? map(t -> deepmap(f, t; filter=filter), tree) : f(tree)
+end
+
 function contraction_sequence(tn::AbstractITensorNetwork; kwargs...)
   # TODO: Use `token_vertex` and/or `token_vertices` here.
   ts = map(v -> tn[v], (1:nv(tn))th)
   seq_linear_index = contraction_sequence(ts; kwargs...)
   # TODO: Use `Functors.fmap` or `StructWalk`?
   return deepmap(n -> Key(vertices(tn)[n * th]), seq_linear_index)
 end
-
-function contraction_sequence(::Algorithm"optimal", tn::Vector{ITensor})
-  return optimal_contraction_sequence(tn)
-end

test/Project.toml

@@ -27,6 +27,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SplitApplyCombine = "03a91e81-4c3e-53e1-a0a4-9c0c8f19dd66"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb"
+TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 UnicodePlots = "b8865327-cd53-5732-bb35-84acbb429228"
 Weave = "44d3d7a6-8a23-5bf8-98c5-b353f8df5ec9"

test/test_additensornetworks.jl

@@ -6,6 +6,7 @@ using ITensorNetworks: ITensorNetwork, inner_network, random_tensornetwork, site
 using ITensors: ITensors, apply, op, scalar, inner
 using LinearAlgebra: norm_sqr
 using StableRNGs: StableRNG
+using TensorOperations: TensorOperations
 using Test: @test, @testset
 @testset "add_itensornetworks" begin
   g = named_grid((2, 2))

test/test_apply.jl

@@ -16,6 +16,7 @@ using NamedGraphs.NamedGraphGenerators: named_grid
 using NamedGraphs.PartitionedGraphs: PartitionVertex
 using SplitApplyCombine: group
 using StableRNGs: StableRNG
+using TensorOperations: TensorOperations
 using Test: @test, @testset
 @testset "apply" begin
   g_dims = (2, 2)

test/test_belief_propagation.jl

@@ -34,6 +34,7 @@ using NamedGraphs.NamedGraphGenerators: named_comb_tree, named_grid
 using NamedGraphs.PartitionedGraphs: PartitionVertex, partitionedges
 using SplitApplyCombine: group
 using StableRNGs: StableRNG
+using TensorOperations: TensorOperations
 using Test: @test, @testset
 
 @testset "belief_propagation (eltype=$elt)" for elt in (