fix sampling for sliced einsum

Author: GiggleLiu

src/sampling.jl

@@ -128,6 +128,13 @@ function sample(tn::TensorNetworkModel, n::Int; usecuda = false)::AbstractMatrix
     return samples.samples
 end
 
+function generate_samples(se::SlicedEinsum, cache::CacheTree{T}, samples, size_dict::Dict) where {T}
+    # slicing is not supported yet.
+    if length(se.slicing) != 0
+        @warn "Slicing is not supported for caching, got nslices = $(length(se.slicing))! Fallback to `NestedEinsum`."
+    end
+    return generate_samples(se.eins, cache, samples, size_dict)
+end
 function generate_samples(code::NestedEinsum, cache::CacheTree{T}, samples, size_dict::Dict) where {T}
     if !OMEinsum.isleaf(code)
         xs = ntuple(i -> cache.siblings[i].content, length(cache.siblings))

test/sampling.jl

@@ -55,7 +55,7 @@ using TensorInference, Test
 
     # fix the evidence
     set_evidence!(instance, 7=>1)
-    tnet = TensorNetworkModel(instance)
+    tnet = TensorNetworkModel(instance, optimizer=TreeSA())
     samples = sample(tnet, n)
     mars = getindex.(marginals(tnet), 1)
     mars_sample = [count(i->samples[k, i]==(0), axes(samples, 2)) for k=1:8] ./ n