save

Author: GiggleLiu

example/asia/asia.jl

@@ -22,6 +22,9 @@ tnet = TensorNetworkModel(instance)
 # Get the maximum log-probabilities (MAP)
 maximum_logp(tnet)
 
+# To sample from the probability model
+sample(tnet, 10)
+
 # Get not only the maximum log-probability, but also the most probable conifguration
 # In the most probable configuration, the most probable one is the patient smoke (3) and has lung cancer (4)
 logp, cfg = most_probable_config(tnet)

src/inference.jl

@@ -49,7 +49,7 @@ function cached_einsum(code::NestedEinsum, @nospecialize(xs), size_dict)
 end
 
 # computed gradient tree by back propagation
-function generate_gradient_tree(se::SlicedEinsum, cache::CacheTree{T}, dy, size_dict::Dict) where {T}
+function generate_gradient_tree(se::SlicedEinsum, cache::CacheTree{T}, dy::AbstractArray{T}, size_dict::Dict) where {T}
     if length(se.slicing) != 0
         @warn "Slicing is not supported for generating masked tree! Fallback to `NestedEinsum`."
     end
@@ -58,7 +58,7 @@ end
 
 # recursively compute the gradients and store it into a tree.
 # also known as the back-propagation algorithm.
-function generate_gradient_tree(code::NestedEinsum, cache::CacheTree{T}, dy, size_dict::Dict) where {T}
+function generate_gradient_tree(code::NestedEinsum, cache::CacheTree{T}, dy::AbstractArray{T}, size_dict::Dict) where {T}
     if OMEinsum.isleaf(code)
         return CacheTree(dy, CacheTree{T}[])
     else

src/sampling.jl

@@ -1,4 +1,14 @@
 ############ Sampling ############
+"""
+$TYPEDEF
+
+### Fields
+$TYPEDFIELDS
+
+The sampled configurations are stored in `samples`, which is a vector of vector.
+`labels` is a vector of variable names for labeling configurations.
+The `setmask` is an boolean indicator to denote whether the sampling process of a variable is complete.
+"""
 struct Samples{L}
     samples::Vector{Vector{Int}}
     labels::Vector{L}
@@ -18,7 +28,7 @@ idx4labels(totalset, labels) = map(v->findfirst(==(v), totalset), labels)
 """
 $(TYPEDSIGNATURES)
 
-The backward rule for tropical einsum.
+The backward process for sampling configurations.
 
 * `ixs` and `xs` are labels and tensor data for input tensors,
 * `iy` and `y` are labels and tensor data for the output tensor,
@@ -72,7 +82,7 @@ function sample(tn::TensorNetworkModel, n::Int; usecuda = false)::Samples
     # forward compute and cache intermediate results.
     cache = cached_einsum(tn.code, xs, size_dict)
     # initialize `y̅` as the initial batch of samples.
-    labels = OMEinsum.uniquelabels(tn.code)
+    labels = get_vars(tn)
     iy = getiyv(tn.code)
     setmask = falses(length(labels))
     idx = map(l->findfirst(==(l), labels), iy)
@@ -86,7 +96,7 @@ function sample(tn::TensorNetworkModel, n::Int; usecuda = false)::Samples
     samples = Samples(configs, labels, setmask)
     # back-propagate
     generate_samples(tn.code, cache, samples, size_dict)
-    return samples
+    return samples.samples
 end
 
 function generate_samples(code::NestedEinsum, cache::CacheTree{T}, samples, size_dict::Dict) where {T}

test/sampling.jl

@@ -50,6 +50,6 @@ using TensorInference, Test
     tnet = TensorNetworkModel(instance)
     samples = sample(tnet, n)
     mars = getindex.(marginals(tnet), 2)
-    mars_sample = [count(s->s[k]==(1), samples.samples) for k=1:8] ./ n
+    mars_sample = [count(s->s[k]==(1), samples) for k=1:8] ./ n
     @test isapprox(mars, mars_sample, atol=0.05)
 end