Public API

Index

Modules

Types

Functions

Modules

TensorInference — Module
source

Types

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

The fast but poor greedy optimizer. Input arguments are

* `α` 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).
* `nrepeat` is the number of repeatition, returns the best contraction order.
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

  • 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

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

  • 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),
  • sub_optimizer, the optimizer for the bipartited sub graphs, one can choose GreedyMethod() or TreeSA(),
  • initializer, the partition configuration initializer, one can choose :random or :greedy (slow but better).

References

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

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 of the unity tensors in the tensors array. Unity tensors are dummy tensors 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 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).
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
  • 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
) -> 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

  • 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