Public API

Index

Modules

Types

Functions

Modules

TensorInference — Module
source

Types

OMEinsumContractionOrders.GreedyMethod — Type
GreedyMethod{MT}
GreedyMethod(; α = 0.0, temperature = 0.0)

It may not be optimal, but it is fast.

Fields

  • α is the parameter for the loss function, for pairwise interaction, L = size(out) - α * (size(in1) + size(in2))
  • temperature is the parameter for sampling, if it is zero, the minimum loss is selected; for non-zero, the loss is selected by the Boltzmann distribution, given by p ~ exp(-loss/temperature).
source
OMEinsumContractionOrders.KaHyParBipartite — Type
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

Fields

  • 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,
  • sub_optimizer is the sub-optimizer used to find the contraction order when the group size is small enough.

References

source
OMEinsumContractionOrders.HyperND — Type
HyperND(;
    dis = KaHyParND(),
    algs = (MF(), AMF(), MMD()),
    level = 6,
    width = 120,
    imbalances = 130:130,
    score = ScoreFunction(),
)

Nested-dissection based optimizer. Recursively partitions a tensor network, then calls a greedy algorithm on the leaves. The optimizer is run a number of times: once for each greedy algorithm in algs and each imbalance value in imbalances. The recursion depth is controlled by the parameters level and width.

The line graph is partitioned using the algorithm dis. OMEinsumContractionOrders currently supports two partitioning algorithms, both of which require importing an external library.

typepackage
METISNDMetis.jl
KaHyParNDKayHyPar.jl

The optimizer is implemented using the tree decomposition library CliqueTrees.jl.

Arguments

source
OMEinsumContractionOrders.TreeSASlicer — Type
TreeSASlicer{IT, LT} <: CodeSlicer

A structure for configuring the Tree Simulated Annealing (TreeSA) slicing algorithm. The goal of slicing is to reach the target space complexity specified by score.sc_target.

Fields

  • 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.
  • fixed_slices::Vector{LT}: A vector of fixed slices that should not be altered. Default is an empty vector.
  • optimization_ratio::Float64: A constant used for determining the number of iterations for slicing. Default is 2.0. i.e. if the current space complexity is 30, and the target space complexity is 20, then the number of iterations for slicing is (30 - 20) x optimization_ratio.
  • score::ScoreFunction: A function to evaluate the quality of the contraction tree. Default is ScoreFunction(sc_target=30.0).
  • decomposition_type::AbstractDecompositionType: The type of decomposition to use. Default is TreeDecomp().

References

source
OMEinsumContractionOrders.ScoreFunction — Type
ScoreFunction

A function to compute the score of a contraction code:

score = tc_weight * 2^tc + rw_weight * 2^rw + sc_weight * max(0, 2^sc - 2^sc_target)

Fields

  • tc_weight: the weight of the time complexity, default is 1.0.
  • sc_weight: the weight of the space complexity (the size of the largest tensor), default is 1.0.
  • rw_weight: the weight of the read-write complexity, default is 0.0.
  • sc_target: the target space complexity, below which the sc_weight will be set to 0 automatically, default is 0.0.
source
OMEinsumContractionOrders.MergeGreedy — Type
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).

source
OMEinsumContractionOrders.SABipartite — Type
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

Fields

  • sc_target is the target space complexity, defined as log2(number of elements in the largest tensor),
  • ntrials is the number of repetition (with different random seeds),
  • βs is a list of inverse temperature 1/T,
  • niters is the number of iteration in each temperature,
  • max_group_size is the maximum size that allowed to used greedy search,
  • sub_optimizer is the optimizer for the bipartited sub graphs, one can choose GreedyMethod() or TreeSA(),
  • initializer is the partition configuration initializer, one can choose :random or :greedy (slow but better).

References

source
OMEinsumContractionOrders.TreeSA — Type
TreeSA{IT, DT} <: CodeOptimizer
TreeSA(; βs=collect(0.01:0.05:15), ntrials=10, niters=50, initializer=:greedy, score=ScoreFunction(), decomposition_type=TreeDeomp())

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

Fields

  • 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.
  • initializer specifies how to determine the initial configuration, it can be :greedy, :random or :specified. If the initializer is :specified, the input code should be a NestedEinsum object.
  • score specifies the score function to evaluate the quality of the contraction tree, it is a function of time complexity, space complexity and read-write complexity.
  • decomposition_type specifies the type of decomposition to use, it can be TreeDeomp or PathDecomp.

References

Breaking changes:

  • nslices is removed, since the slicing part is now separated from the optimization part, see slice_code function and TreeSASlicer.
  • greedy_method is removed. If you want to have detailed control of the initializer, please pre-optimize the code with another method and then use :specified to initialize the tree.
source
TensorInference.MMAPModel — Type
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.
source
TensorInference.RescaledArray — Type
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.

source
TensorInference.TensorNetworkModel — Type
struct TensorNetworkModel{ET, MT<:AbstractArray}

Probabilistic modeling with a tensor network.

Fields

  • nvars are the number of variables 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 unity_tensors_labels.
  • evidence is a dictionary used to specify degrees of freedom that are fixed to certain values.
  • unity_tensors_idx is a vector of indices pointing to the unity tensors in the tensors array. Unity tensors are dummy tensors with all entries equal to one, which are used to obtain the marginal probabilities.
source
TensorInference.UAIModel — Type
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,
source
TensorInference.BeliefPropgation — Type
struct BeliefPropgation{T}
BeliefPropgation(nvars::Int, t2v::AbstractVector{Vector{Int}}, tensors::AbstractVector{AbstractArray{T}}) where T

A belief propagation object.

Fields

  • t2v::Vector{Vector{Int}}: a mapping from tensors to variables
  • v2t::Vector{Vector{Int}}: a mapping from variables to tensors
  • tensors::Vector{AbstractArray{T}}: the tensors
source

Functions

OMEinsumContractionOrders.contraction_complexity — Function
contraction_complexity(tensor_network)

Returns the contraction complexity of a tensor newtork model.

source
contraction_complexity(eincode, size_dict) -> ContractionComplexity

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

  • tc: time complexity defined as log2(number of element-wise multiplications).
  • sc: space complexity defined as log2(size of the maximum intermediate tensor).
  • rwc: read-write complexity defined as log2(the number of read-write operations).
source
TensorInference.get_cards — Function
get_cards(tn::TensorNetworkModel; fixedisone) -> Vector

Get the ardinalities of variables in this tensor network.

source
get_cards(mmap::MMAPModel; fixedisone) -> Vector
source
TensorInference.get_vars — Function
get_vars(tn::TensorNetworkModel) -> Vector{Int64}

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

source
get_vars(mmap::MMAPModel) -> Vector
source
TensorInference.log_probability — Function
log_probability(
    tn::TensorNetworkModel,
    config::Union{Dict, AbstractVector}
) -> Real

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

source
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.

source
TensorInference.marginals — Function
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 unity_tensors_labels field when constructing a TensorNetworkModel. Note that the choice of marginal variables will affect the contraction order of the tensor network.

Arguments

Keyword Arguments

  • usecuda::Bool: Specifies whether to use CUDA for tensor contraction.
  • rescale::Bool: 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), unity_tensors_labels = [[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.

source
marginals(
    state::TensorInference.BPState{T, VT} where VT<:AbstractArray{T, 1}
) -> Dict
source
TensorInference.maximum_logp — Function
maximum_logp(
    tn::TensorNetworkModel;
    usecuda
) -> AbstractArray{<:Real}

Returns an output array containing largest log-probabilities.

source
TensorInference.probability — Function
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.

source
TensorInference.belief_propagate — Function
belief_propagate(
    bp::BeliefPropgation;
    kwargs...
) -> Tuple{TensorInference.BPState{T, VT} where {T, VT<:Vector{T}}, TensorInference.BPInfo}

Run the belief propagation algorithm, and return the final state and the information about the convergence.

Arguments

  • bp::BeliefPropgation: the belief propagation object

Keyword Arguments

  • max_iter::Int=100: the maximum number of iterations
  • tol::Float64=1e-6: the tolerance for the convergence, the convergence is checked by infidelity of messages in consecutive iterations. For complex numbers, the converged message may be different only by a phase factor.
  • damping::Float64=0.2: the damping factor for the message update, updated-message = damping * old-message + (1 - damping) * new-message
source
TensorInference.dataset_from_artifact — Function
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.

source
TensorInference.problem_from_artifact — Function
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.

source
TensorInference.read_solution — Function
read_solution(
    problem::ArtifactProblemSpec;
    factor_eltype
) -> Union{Nothing, Float64, Vector}

Return the solution in the artifact.

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

source
TensorInference.read_model_file — Function
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/

source
TensorInference.read_evidence_file — Function
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/

source
TensorInference.read_td_file — Function
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/

source
TensorInference.sample — Function
sample(
    tn::TensorNetworkModel,
    n::Int64;
    usecuda,
    queryvars,
    rescale
) -> TensorInference.Samples{Int64}

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

  • rescale is a boolean flag to indicate whether to rescale the tensors during contraction.
  • 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).
source
TensorInference.update_evidence! — Function
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.
source
TensorInference.update_temperature — Function
update_temperature(
    tnet::TensorNetworkModel,
    problem::ProblemReductions.ConstraintSatisfactionProblem,
    β::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.
source
TensorInference.random_matrix_product_state — Function
random_matrix_product_state(
    ::Type{T},
    n::Int64,
    chi::Int64
) -> TensorNetworkModel{OMEinsum.DynamicNestedEinsum{Int64}}
random_matrix_product_state(
    ::Type{T},
    n::Int64,
    chi::Int64,
    d::Int64
) -> TensorNetworkModel{OMEinsum.DynamicNestedEinsum{Int64}}

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.
source
TensorInference.random_tensor_train_uai — Function
random_tensor_train_uai(
    ::Type{T},
    n::Int64,
    chi::Int64;
    ...
) -> UAIModel
random_tensor_train_uai(
    ::Type{T},
    n::Int64,
    chi::Int64,
    d::Int64;
    periodic
) -> UAIModel

Tensor train (TT) is a tensor network model that is widely used in quantum many-body physics. This model is different from the matrix product state (MPS) in that it does not have an extra copy for representing the bra state.

source
TensorInference.save_tensor_network — Function
save_tensor_network(tn::TensorNetworkModel; folder::String)

Save a tensor network model to a folder with separate files for code, tensors, and model metadata. The code is saved using OMEinsum.writejson, tensors as JSON, and model specifics in model.json.

Arguments

  • tn::TensorNetworkModel: The tensor network model to save
  • folder::String: The folder path to save the files

Files Created

  • code.json: Contains the einsum code using OMEinsum format
  • tensors.json: Contains the tensor data as JSON
  • model.json: Contains nvars, evidence, and unitytensorsidx

Example

tn = TensorNetworkModel(...)  # create your model
save_tensor_network(tn; folder="my_model")
source
TensorInference.load_tensor_network — Function
load_tensor_network(folder::String)

Load a tensor network model from a folder containing code, tensors, and model files.

Arguments

  • folder::String: The folder path containing the files

Returns

  • TensorNetworkModel: The loaded tensor network model

Required Files

  • code.json: Contains the einsum code using OMEinsum format
  • tensors.json: Contains the tensor data as JSON
  • model.json: Contains nvars, evidence, and unitytensorsidx

Example

tn = load_tensor_network("my_model")
source