Contraction sequence optimization

When contracting a tensor network, the sequence of contraction makes a big difference in the computational cost. However, the complexity of determining the optimal sequence grows exponentially with the number of tensors, but there are many heuristic algorithms available for computing optimal sequences for small networks^[1]^[2]^[3]^[4]^[5]^[6]. ITensors.jl imports functionality from TensorOperations.jl for helping you find the optimal contraction sequence for small tensor network, as we will show below.

Functions

ITensors.optimal_contraction_sequence — Function

optimal_contraction_sequence(T)

Returns a contraction sequence for contracting the tensors T. The sequence is generally optimal and is found via the optimaltree function in TensorOperations.jl which must be loaded.

source

ITensors.contraction_cost — Function

contraction_cost(A; sequence)

Return the cost of contracting the collection of ITensors according to the specified sequence, where the cost is measured in the number of floating point operations that would need to be performed to contract dense tensors of the dimensions specified by the indices of the tensors (so for now, sparsity is ignored in computing the costs). Pairwise costs are returned in a vector (contracting N tensors requires N-1 pairwise contractions). You can use sum(contraction_cost(A; sequence)) to get the total cost of the contraction.

If no sequence is specified, left associative contraction is used, in other words the sequence is equivalent to [[[[1, 2], 3], 4], …].

source

NDTensors.contract — Function

*(As::ITensor...; sequence = default_sequence(), kwargs...)
*(As::Vector{<: ITensor}; sequence = default_sequence(), kwargs...)
contract(As::ITensor...; sequence = default_sequence(), kwargs...)

Contract the set of ITensors according to the contraction sequence.

The default sequence is "automatic" if ITensors.using_contraction_sequence_optimization() is true, otherwise it is "left_associative" (the ITensors are contracted from left to right).

You can change the default with ITensors.enable_contraction_sequence_optimization() and ITensors.disable_contraction_sequence_optimization().

For a custom sequence, the sequence should be provided as a binary tree where the leaves are integers n specifying the ITensor As[n] and branches are accessed by indexing with 1 or 2, i.e. sequence = Any[Any[1, 3], Any[2, 4]].

source

Examples

In the following example we show how to compute the contraction sequence cost of a

using ITensors
using Symbolics

using ITensors: contraction_cost

@variables m, k, d

l = Index(m, "l")
r = Index(m, "r")
h₁ = Index(k, "h₁")
h₂ = Index(k, "h₂")
h₃ = Index(k, "h₃")
s₁ = Index(d, "s₁")
s₂ = Index(d, "s₂")

H₁ = ITensor(dag(s₁), s₁', dag(h₁), h₂)
H₂ = ITensor(dag(s₂), s₂', dag(h₂), h₃)
L = ITensor(dag(l), l', h₁)
R = ITensor(dag(r), r', h₃)
ψ = ITensor(l, s₁, s₂, r)

TN = [ψ, L, H₁, H₂, R]
sequence1 = Any[2, Any[3, Any[4, Any[1, 5]]]]
sequence2 = Any[Any[4, 5], Any[1, Any[2, 3]]]
cost1 = contraction_cost(TN; sequence = sequence1)
cost2 = contraction_cost(TN; sequence = sequence2)

println("First sequence")
display(sequence1)
display(cost1)
@show sum(cost1)
@show substitute(sum(cost1), Dict(d => 4))

println("\nSecond sequence")
display(sequence2)
display(cost2)
@show sum(cost2)
@show substitute(sum(cost2), Dict(d => 4))

This example helps us learn that in the limit of large MPS bond dimension m, the first contraction sequence is faster, while in the limit of large MPO bond dimension k, the second sequence is faster. This has practical implications for writing an efficient DMRG algorithm in both limits, which we plan to incorporate into ITensors.jl.

Here is a more systematic example of searching through the parameter space to find optimal contraction sequences. Note, the TensorOperations.jl library must be loaded to use the optimalcontractionsequence function:

using ITensors
using Symbolics

using ITensors: contraction_cost, optimal_contraction_sequence
using TensorOperations: TensorOperations

function tensor_network(; m, k, d)
  l = Index(m, "l")
  r = Index(m, "r")
  h₁ = Index(k, "h₁")
  h₂ = Index(k, "h₂")
  h₃ = Index(k, "h₃")
  s₁ = Index(d, "s₁")
  s₂ = Index(d, "s₂")

  ψ = ITensor(l, s₁, s₂, r)
  L = ITensor(dag(l), l', h₁)
  H₁ = ITensor(dag(s₁), s₁', dag(h₁), h₂)
  H₂ = ITensor(dag(s₂), s₂', dag(h₂), h₃)
  R = ITensor(dag(r), r', h₃)
  return [ψ, L, H₁, H₂, R]
end

function main()
  mrange = 50:10:80
  krange = 50:10:80
  sequence_costs = Matrix{Any}(undef, length(mrange), length(krange))
  for iₘ in eachindex(mrange), iₖ in eachindex(krange)
    m_val = mrange[iₘ]
    k_val = krange[iₖ]
    d_val = 4

    TN = tensor_network(; m = m_val, k = k_val, d = d_val)
    sequence = optimal_contraction_sequence(TN)
    cost = contraction_cost(TN; sequence = sequence)

    @variables m, k, d
    TN_symbolic = tensor_network(; m = m, k = k, d = d)
    cost_symbolic = contraction_cost(TN_symbolic; sequence = sequence)
    sequence_cost = (dims = (m = m_val, k = k_val, d = d_val), sequence = sequence, cost = cost, symbolic_cost = cost_symbolic)
    sequence_costs[iₘ, iₖ] = sequence_cost
  end
  return sequence_costs
end

sequence_costs = main()

# Analyze the results.
println("Index dimensions")
display(getindex.(sequence_costs, :dims))

println("\nContraction sequences")
display(getindex.(sequence_costs, :sequence))

println("\nSymbolic contraction cost with d = 4")
# Fix d to a certain value (such as 4 for a Hubbard site)
@variables d
var_sub = Dict(d => 4)
display(substitute.(sum.(getindex.(sequence_costs, :symbolic_cost)), (var_sub,)))

A future direction will be to allow optimizing over contraction sequences with the dimensions specified symbolically, so that the optimal sequence in limits of certain dimensions can be found. In addition, we plan to implement more algorithms that work for larger networks, as well as algorithms like^[2] which take an optimal sequence for a closed network and generate optimal sequences for environments of each tensor in the network, which is helpful for computing gradients of tensor networks.