Public API

Main module for TensorInference.jl – A toolbox for probabilistic inference using contraction of tensor networks.

Exports

• ArtifactProblemSpec
• GreedyMethod
• KaHyParBipartite
• MMAPModel
• MergeGreedy
• MergeVectors
• RescaledArray
• SABipartite
• TensorNetworkModel
• TreeSA
• UAIModel
• contraction_complexity
• dataset_from_artifact
• get_cards
• get_vars
• log_probability
• marginals
• maximum_logp
• most_probable_config
• probability
• problem_from_artifact
• random_matrix_product_state
• read_evidence
• read_evidence_file
• read_model
• read_model_file
• read_queryvars
• read_solution
• read_td_file
• sample
• update_evidence!
• update_temperature

GreedyMethod{MT}
GreedyMethod(; method=MinSpaceOut(), nrepeat=10)

The fast but poor greedy optimizer. Input arguments are

• method is MinSpaceDiff() or MinSpaceOut.
  • MinSpaceOut choose one of the contraction that produces a minimum output tensor size,
  • MinSpaceDiff choose one of the contraction that decrease the space most.
• nrepeat is the number of repeatition, returns the best contraction order.

KaHyParBipartite{RT,IT,GM}
KaHyParBipartite(; sc_target, imbalances=collect(0.0:0.005:0.8),
    max_group_size=40, greedy_config=GreedyMethod())

Optimize the einsum code contraction order using the KaHyPar + Greedy approach. This program first recursively cuts the tensors into several groups using KaHyPar, with maximum group size specifed by max_group_size and maximum space complexity specified by sc_target, Then finds the contraction order inside each group with the greedy search algorithm. Other arguments are

• sc_target is the target space complexity, defined as log2(number of elements in the largest tensor),
• imbalances is a KaHyPar parameter that controls the group sizes in hierarchical bipartition,
• max_group_size is the maximum size that allowed to used greedy search,
• greedy_config is a greedy optimizer.

References

• Hyper-optimized tensor network contraction
• Simulating the Sycamore quantum supremacy circuits

MergeGreedy <: CodeSimplifier
MergeGreedy(; threshhold=-1e-12)

Contraction code simplifier (in order to reduce the time of calling optimizers) that merges tensors greedily if the space complexity of merged tensors is reduced (difference smaller than the threshhold).

MergeVectors <: CodeSimplifier
MergeVectors()

Contraction code simplifier (in order to reduce the time of calling optimizers) that merges vectors to closest tensors.

SABipartite{RT,BT}
SABipartite(; sc_target=25, ntrials=50, βs=0.1:0.2:15.0, niters=1000
    max_group_size=40, greedy_config=GreedyMethod(), initializer=:random)

Optimize the einsum code contraction order using the Simulated Annealing bipartition + Greedy approach. This program first recursively cuts the tensors into several groups using simulated annealing, with maximum group size specifed by max_group_size and maximum space complexity specified by sc_target, Then finds the contraction order inside each group with the greedy search algorithm. Other arguments are

• size_dict, a dictionary that specifies leg dimensions,
• sc_target is the target space complexity, defined as log2(number of elements in the largest tensor),
• max_group_size is the maximum size that allowed to used greedy search,
• βs is a list of inverse temperature 1/T,
• niters is the number of iteration in each temperature,
• ntrials is the number of repetition (with different random seeds),
• greedy_config configures the greedy method,
• initializer, the partition configuration initializer, one can choose :random or :greedy (slow but better).

References

• Hyper-optimized tensor network contraction

TreeSA{RT,IT,GM}
TreeSA(; sc_target=20, βs=collect(0.01:0.05:15), ntrials=10, niters=50,
    sc_weight=1.0, rw_weight=0.2, initializer=:greedy, greedy_config=GreedyMethod(; nrepeat=1))

Optimize the einsum contraction pattern using the simulated annealing on tensor expression tree.

• sc_target is the target space complexity,
• ntrials, βs and niters are annealing parameters, doing ntrials indepedent annealings, each has inverse tempteratures specified by βs, in each temperature, do niters updates of the tree.
• sc_weight is the relative importance factor of space complexity in the loss compared with the time complexity.
• rw_weight is the relative importance factor of memory read and write in the loss compared with the time complexity.
• initializer specifies how to determine the initial configuration, it can be :greedy or :random. If it is using :greedy method to generate the initial configuration, it also uses two extra arguments greedy_method and greedy_nrepeat.
• nslices is the number of sliced legs, default is 0.
• fixed_slices is a vector of sliced legs, default is [].

References

• Recursive Multi-Tensor Contraction for XEB Verification of Quantum Circuits

struct MMAPModel{LT, AT<:AbstractArray}

Computing the most likely assignment to the query variables, Xₘ ⊆ X after marginalizing out the remaining variables Xₛ = X \ Xₘ.

\[{\rm MMAP}(X_i|E=e) = \arg \max_{X_M} \sum_{X_S} \prod_{F} f(x_M, x_S, e)\]

Fields

• vars is the query variables in the tensor network.
• code is the tropical tensor network contraction pattern.
• tensors is the tensors fed into the tensor network.
• clusters is the clusters, each element of this cluster is a TensorNetworkModel instance for marginalizing certain variables.
• evidence is a dictionary to specifiy degree of freedoms fixed to certain values, which should not have overlap with the query variables.

struct RescaledArray{T, N, AT<:AbstractArray{T, N}} <: AbstractArray{T, N}

RescaledArray(α, T) -> RescaledArray

An array data type with a log-prefactor, and a l∞-normalized storage, i.e. the maximum element in a tensor is 1. This tensor type can avoid the potential underflow/overflow of numbers in a tensor network. The constructor RescaledArray(α, T) creates a rescaled array that equal to exp(α) * T.

struct TensorNetworkModel{LT, ET, MT<:AbstractArray}

Probabilistic modeling with a tensor network.

Fields

• vars are the degrees of freedom in the tensor network.
• code is the tensor network contraction pattern.
• tensors are the tensors fed into the tensor network, the leading tensors are unity tensors associated with mars.
• evidence is a dictionary used to specify degrees of freedom that are fixed to certain values.
• mars is a vector, each element is a vector of variables to compute marginal probabilities.

struct ArtifactProblemSpec

Specify the UAI models from the artifacts. It can be used as the input of read_model.

Fields

• artifact_path::String

• task::String

• problem_set::String

• problem_id::Int64

struct UAIModel{ET, FT<:(TensorInference.Factor{ET})}

Fields

• nvars is the number of variables,
• cards is a vector of cardinalities for variables,
• factors is a vector of factors,

contraction_complexity(tensor_network)

Returns the contraction complexity of a tensor newtork model.

contraction_complexity(eincode, size_dict) -> ContractionComplexity

Returns the time, space and read-write complexity of the einsum contraction. The returned object contains 3 fields:

• time complexity tc defined as log2(number of element-wise multiplications).
• space complexity sc defined as log2(size of the maximum intermediate tensor).
• read-write complexity rwc defined as log2(the number of read-write operations).

get_cards(tn::TensorNetworkModel; fixedisone) -> Vector

Get the cardinalities of variables in this tensor network.

get_cards(mmap::MMAPModel; fixedisone) -> Vector

get_vars(tn::TensorNetworkModel) -> Vector

Get the variables in this tensor network, they are also known as legs, labels, or degree of freedoms.

get_vars(mmap::MMAPModel) -> Vector

log_probability(
    tn::TensorNetworkModel,
    config::Union{Dict, AbstractVector}
) -> Real

Evaluate the log probability (or partition function) of config.

log_probability(
    tn::TensorNetworkModel;
    usecuda
) -> AbstractArray

Evaluate the log probability (or partition function). It is the logged version of probability, which is less likely to overflow.

marginals(
    tn::TensorNetworkModel;
    usecuda,
    rescale
) -> Dict{Vector{Int64}}

Queries the marginals of the variables in a TensorNetworkModel. The function returns a dictionary, where the keys are the variables and the values are their respective marginals. A marginal is a probability distribution over a subset of variables, obtained by integrating or summing over the remaining variables in the model. By default, the function returns the marginals of all individual variables. To specify which marginal variables to query, set the mars field when constructing a TensorNetworkModel. Note that the choice of marginal variables will affect the contraction order of the tensor network.

Arguments

• tn: The TensorNetworkModel to query.
• usecuda: Specifies whether to use CUDA for tensor contraction.
• rescale: Specifies whether to rescale the tensors during contraction.

Example

The following example is taken from examples/asia-network/main.jl.

julia> model = read_model_file(pkgdir(TensorInference, "examples", "asia-network", "model.uai"));

julia> tn = TensorNetworkModel(model; evidence=Dict(1=>0));

julia> marginals(tn)
Dict{Vector{Int64}, Vector{Float64}} with 8 entries:
  [8] => [0.450138, 0.549863]
  [3] => [0.5, 0.5]
  [1] => [1.0]
  [5] => [0.45, 0.55]
  [4] => [0.055, 0.945]
  [6] => [0.10225, 0.89775]
  [7] => [0.145092, 0.854908]
  [2] => [0.05, 0.95]

julia> tn2 = TensorNetworkModel(model; evidence=Dict(1=>0), mars=[[2, 3], [3, 4]]);

julia> marginals(tn2)
Dict{Vector{Int64}, Matrix{Float64}} with 2 entries:
  [2, 3] => [0.025 0.025; 0.475 0.475]
  [3, 4] => [0.05 0.45; 0.005 0.495]

In this example, we first set the evidence for variable 1 to 0 and then query the marginals of all individual variables. The returned dictionary has keys that correspond to the queried variables and values that represent their marginals. These marginals are vectors, with each entry corresponding to the probability of the variable taking a specific value. In this example, the possible values are 0 or 1. For the evidence variable 1, the marginal is always [1.0] since its value is fixed at 0.

Next, we specify the marginal variables to query as variables 2 and 3, and variables 3 and 4, respectively. The joint marginals may or may not affect the contraction time and space. In this example, the contraction space complexity increases from 2^{2.0} to 2^{5.0}, and the contraction time complexity increases from 2^{5.977} to 2^{7.781}. The output marginals are the joint probabilities of the queried variables, represented by tensors.

maximum_logp(
    tn::TensorNetworkModel;
    usecuda
) -> AbstractArray{<:Real}

Returns an output array containing largest log-probabilities.

most_probable_config(
    tn::TensorNetworkModel;
    usecuda
) -> Tuple{Real, Vector}

Returns the largest log-probability and the most probable configuration.

probability(
    tn::TensorNetworkModel;
    usecuda,
    rescale
) -> AbstractArray

Contract the tensor network and return an array of probability of evidence. Precisely speaking, the return value is the partition function, which may not be l1-normalized.

If the openvars of the input tensor networks is zero, the array rank is zero. Otherwise, the return values corresponds to marginal probabilities.

dataset_from_artifact(
    artifact_name::AbstractString
) -> Dict{String, Dict{String, Dict{Int64, ArtifactProblemSpec}}}

Helper function that captures the problem names that belong to problem_set for the given task.

problem_from_artifact(
    artifact_name::String,
    task::String,
    problem_set::String,
    problem_id::Int64
) -> ArtifactProblemSpec

Get artifact from artifact name, task name, problem set name and problem id.

read_model(problem::ArtifactProblemSpec; eltype) -> UAIModel

Read an UAI model from an artifact.

read_evidence(
    problem::ArtifactProblemSpec
) -> Dict{Int64, Int64}

read_solution(
    problem::ArtifactProblemSpec;
    factor_eltype
) -> Any

Return the solution in the artifact.

The UAI file formats are defined in: https://uaicompetition.github.io/uci-2022/file-formats/

read_queryvars(
    problem::ArtifactProblemSpec
) -> Vector{Int64}

read_model_file(
    model_filepath::AbstractString;
    factor_eltype
) -> UAIModel

Parse the problem instance found in model_filepath defined in the UAI model format. If the provided file path is empty, return nothing.

The UAI file formats are defined in: https://uaicompetition.github.io/uci-2022/file-formats/

read_evidence_file(
    evidence_filepath::AbstractString
) -> Tuple{Vector{Int64}, Vector{Int64}}

Return the observed variables and values in evidence_filepath. If the passed file path is an empty string, return empty vectors.

The UAI file formats are defined in: https://uaicompetition.github.io/uci-2022/file-formats/

read_td_file(
    td_filepath::AbstractString
) -> Tuple{Int64, Int64, Int64, Vector{Vector{Int64}}, Vector{Vector{Int64}}}

Parse a tree decomposition instance described the PACE format.

The PACE file format is defined in: https://pacechallenge.org/2017/treewidth/

sample(
    tn::TensorNetworkModel,
    n::Int64;
    usecuda,
    queryvars
) -> TensorInference.Samples

Generate samples from a tensor network based probabilistic model. Returns a vector of vector, each element being a configurations defined on get_vars(tn).

Arguments

• tn is the tensor network model.
• n is the number of samples to be returned.

Keyword Arguments

• usecuda is a boolean flag to indicate whether to use CUDA for tensor computation.
• queryvars is the variables to be sampled, default is get_vars(tn).

update_evidence!(
    tnet::TensorNetworkModel,
    evidence::Dict
) -> TensorNetworkModel

Update the evidence of a tensor network model, without changing the set of observed variables!

Arguments

• tnet is the TensorNetworkModel instance.
• evidence is the new evidence, the keys must be a subset of existing evidence.

update_temperature(
    tnet::TensorNetworkModel,
    problem::GenericTensorNetworks.GraphProblem,
    β::Real
) -> TensorNetworkModel

Update the temperature of a tensor network model. The program will regenerate tensors from the problem, without repeated optimizing the contraction order.

Arguments

• tnet is the TensorNetworkModel instance.
• problem is the target constraint satisfiability problem.
• β is the inverse temperature.

random_matrix_product_state(
    ::Type{T},
    n::Int64,
    chi::Int64
) -> TensorNetworkModel{Int64, ET} where ET<:OMEinsum.DynamicNestedEinsum
random_matrix_product_state(
    ::Type{T},
    n::Int64,
    chi::Int64,
    d::Int64
) -> TensorNetworkModel{Int64, ET} where ET<:OMEinsum.DynamicNestedEinsum

Matrix product state (MPS) is a tensor network model that is widely used in quantum many-body physics. It is a special case of tensor network model where the tensors are rank-3 tensors and the physical indices are connected in a chain. The MPS is defined as:

\[\begin{align*} \left| \psi \right\rangle &= \sum_{x_1, x_2, \ldots, x_n} \text{Tr}(A_1^{x_1} A_2^{x_2} \cdots A_n^{x_n}) \left| x_1, x_2, \ldots, x_n \right\rangle \\ \left\langle \psi \right| &= \sum_{x_1, x_2, \ldots, x_n} \text{Tr}(A_n^{x_n} \cdots A_2^{x_2} A_1^{x_1}) \left\langle x_1, x_2, \ldots, x_n \right| \end{align*}\]

where $A_i^{x_i}$ is a rank-3 tensor with physical index $x_i$ and two virtual indices connecting to the next tensor. The MPS is a special case of the tensor network model where the tensors are rank-3 tensors and the physical indices are connected in a chain.

Arguments

• n is the number of physical indices.
• chi is the bond dimension of the virtual indices.
• d is the dimension of the physical indices.