From 36e90d0849890399182a11e9eacbcaa1ce2448b2 Mon Sep 17 00:00:00 2001
From: GiggleLiu <cacate0129@gmail.com>
Date: Thu, 6 Oct 2022 14:59:14 +0800
Subject: [PATCH 1/6] fix rescale

---
 src/Core.jl            | 10 +++++-----
 src/TensorInference.jl |  2 +-
 src/inference.jl       | 14 +++++++++++++-
 test/inference.jl      | 27 +++++++++++++++------------
 4 files changed, 34 insertions(+), 19 deletions(-)

diff --git a/src/Core.jl b/src/Core.jl
index 09d2ef7..f46e2a9 100644
--- a/src/Core.jl
+++ b/src/Core.jl
@@ -83,21 +83,21 @@ end
 """
 $(TYPEDSIGNATURES)
 """
-function TensorNetworkModeling(instance::UAIInstance; openvertices=(), optimizer=GreedyMethod(), simplifier=nothing)::TensorNetworkModeling
-    return TensorNetworkModeling(1:instance.nvars, instance.factors; fixedvertices=Dict(zip(instance.obsvars, instance.obsvals .- 1)), optimizer, simplifier, openvertices)
+function TensorNetworkModeling(instance::UAIInstance; rescale=1.0, openvertices=(), optimizer=GreedyMethod(), simplifier=nothing)::TensorNetworkModeling
+    return TensorNetworkModeling(1:instance.nvars, instance.cards, instance.factors; fixedvertices=Dict(zip(instance.obsvars, instance.obsvals .- 1)), optimizer, simplifier, openvertices, rescale)
 end
 
 """
 $(TYPEDSIGNATURES)
 """
-function TensorNetworkModeling(vars::AbstractVector{LT}, factors::Vector{<:Factor{T}}; openvertices=(), fixedvertices=Dict{LT,Int}(), optimizer=GreedyMethod(), simplifier=nothing)::TensorNetworkModeling where {T,LT}
+function TensorNetworkModeling(vars::AbstractVector{LT}, cards::AbstractVector{Int}, factors::Vector{<:Factor{T}}; rescale=1.0, openvertices=(), fixedvertices=Dict{LT,Int}(), optimizer=GreedyMethod(), simplifier=nothing)::TensorNetworkModeling where {T,LT}
     # The 1st argument of `EinCode` is a vector of vector of labels for specifying the input tensors, 
     # The 2nd argument of `EinCode` is a vector of labels for specifying the output tensor,
     # e.g.
     # `EinCode([[1, 2], [2, 3]], [1, 3])` is the EinCode for matrix multiplication.
     rawcode = EinCode([[[var] for var in vars]..., [[factor.vars...] for factor in factors]...], collect(LT, openvertices))  # labels for vertex tensors (unity tensors) and edge tensors
-    tensors = [[ones(T, 2) for _=1:length(vars)]..., getfield.(factors, :vals)...]
-    return TensorNetworkModeling(collect(LT, vars), rawcode, tensors; fixedvertices, optimizer, simplifier)
+    tensors = [[ones(T, cards[i]) for i=1:length(vars)]..., getfield.(factors, :vals)...]
+    return TensorNetworkModeling(collect(LT, vars), rawcode, map(t->t .* rescale, tensors); fixedvertices, optimizer, simplifier)
 end
 """
 $(TYPEDSIGNATURES)
diff --git a/src/TensorInference.jl b/src/TensorInference.jl
index 8487a3f..95d9543 100644
--- a/src/TensorInference.jl
+++ b/src/TensorInference.jl
@@ -11,7 +11,7 @@ export timespace_complexity, timespacereadwrite_complexity, TreeSA, GreedyMethod
 export read_uai_file, read_td_file, read_uai_evid_file, read_uai_mar_file, read_uai_problem
 
 # marginals
-export TensorNetworkModeling, get_vars, get_cards, probability, marginals
+export TensorNetworkModeling, get_vars, get_cards, probability, marginals, autorescale!
 
 # MAP
 export most_probable_config, maximum_logp
diff --git a/src/inference.jl b/src/inference.jl
index 1cdd31c..e0d5031 100644
--- a/src/inference.jl
+++ b/src/inference.jl
@@ -127,6 +127,18 @@ One can use `get_vars(tn)` to get the full list of variables in this tensor netw
 """
 function marginals(tn::TensorNetworkModeling; usecuda=false)::Vector
     vars = get_vars(tn)
-    _, grads = cost_and_gradient(tn.code, generate_tensors(tn; usecuda))
+    # sometimes, the cost can overflow.
+    cost, grads = cost_and_gradient(tn.code, generate_tensors(tn; usecuda))
     return LinearAlgebra.normalize!.(grads[1:length(vars)], 1)
+end
+
+"""
+$(TYPEDSIGNATURES)
+"""
+function autorescale!(tn::TensorNetworkModeling; usecuda::Bool=false)::TensorNetworkModeling
+    vars = get_vars(tn)
+    for t in tn.tensors[length(vars)+1:end]
+        t ./= maximum(t)/sqrt(length(t))
+    end
+    return tn
 end
\ No newline at end of file
diff --git a/test/inference.jl b/test/inference.jl
index f5d6399..48fe19b 100644
--- a/test/inference.jl
+++ b/test/inference.jl
@@ -6,20 +6,22 @@ using TensorInference
 
   @testset "UAI 2014 problem set" begin
 
+    # TODO: rescale while contract.
     benchmarks = [
-                  "Alchemy", # fails
-                  "CSP",
-                  "DBN",
-                  "Grids",
-                  "linkage",
-                  "ObjectDetection",
-                  "Pedigree",
-                  "Promedus",
-                  "relational",
-                  "Segmentation",
+                  ("Alchemy", 1/3), # fails
+                  ("CSP",1.0),
+                  #("DBN",1.0),
+                  ("Grids",1/2),
+                  ("linkage",1/2),
+                  ("ObjectDetection",1/2),
+                  #("Pedigree",1.0),
+                  ("Promedus",1.0),
+                  #("relational",1.0),
+                  ("Segmentation",1.0),
                  ]
+    benchmarks = [("linkage",1/3)]
 
-    for benchmark in benchmarks
+    for (benchmark, rescale) in benchmarks
 
       @testset "$(benchmark) benchmark" begin
 
@@ -37,7 +39,8 @@ using TensorInference
             problem = read_uai_problem(problem)
 
             # does not optimize over open vertices
-            tn = TensorNetworkModeling(problem; optimizer=TreeSA(ntrials=1, niters=2, βs=1:0.1:40))
+            tn = TensorNetworkModeling(problem; optimizer=TreeSA(ntrials=3, niters=5, βs=1:0.1:40))
+            autorescale!(tn)
             # @info(tn)
             # @info timespace_complexity(tn)
             marginals2 = marginals(tn)

From c9ffd4c7989aa5a3d3c00ce4fbce9b21d197e0fb Mon Sep 17 00:00:00 2001
From: GiggleLiu <cacate0129@gmail.com>
Date: Sat, 8 Oct 2022 00:19:48 +0800
Subject: [PATCH 2/6] rescale and choose the contraction order finding

---
 Project.toml           |  3 ++-
 src/Core.jl            | 15 ++++++++-------
 src/TensorInference.jl |  2 +-
 src/inference.jl       | 12 +-----------
 src/utils.jl           |  8 +++-----
 test/inference.jl      | 41 ++++++++++++++++++++++++-----------------
 6 files changed, 39 insertions(+), 42 deletions(-)

diff --git a/Project.toml b/Project.toml
index e8d153e..b07b90d 100644
--- a/Project.toml
+++ b/Project.toml
@@ -22,7 +22,8 @@ julia = "1.3"
 
 [extras]
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+KaHyPar = "2a6221f6-aa48-11e9-3542-2d9e0ef01880"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "Documenter"]
+test = ["Test", "Documenter", "KaHyPar"]
diff --git a/src/Core.jl b/src/Core.jl
index f46e2a9..ca3341c 100644
--- a/src/Core.jl
+++ b/src/Core.jl
@@ -96,8 +96,8 @@ function TensorNetworkModeling(vars::AbstractVector{LT}, cards::AbstractVector{I
     # e.g.
     # `EinCode([[1, 2], [2, 3]], [1, 3])` is the EinCode for matrix multiplication.
     rawcode = EinCode([[[var] for var in vars]..., [[factor.vars...] for factor in factors]...], collect(LT, openvertices))  # labels for vertex tensors (unity tensors) and edge tensors
-    tensors = [[ones(T, cards[i]) for i=1:length(vars)]..., getfield.(factors, :vals)...]
-    return TensorNetworkModeling(collect(LT, vars), rawcode, map(t->t .* rescale, tensors); fixedvertices, optimizer, simplifier)
+    tensors = Array{T}[[ones(T, cards[i]) for i=1:length(vars)]..., [t.vals .* rescale for t in factors]...]
+    return TensorNetworkModeling(collect(LT, vars), rawcode, tensors; fixedvertices, optimizer, simplifier)
 end
 """
 $(TYPEDSIGNATURES)
@@ -107,7 +107,9 @@ function TensorNetworkModeling(vars::AbstractVector{LT}, rawcode::EinCode, tenso
     # The 1st argument is the contraction pattern to be optimized (without contraction order).
     # The 2nd arugment is the size dictionary, which is a label-integer dictionary.
     # The 3rd and 4th arguments are the optimizer and simplifier that configures which algorithm to use and simplify.
-    code = optimize_code(rawcode, OMEinsum.get_size_dict(getixsv(rawcode), tensors), optimizer, simplifier)
+    size_dict = OMEinsum.get_size_dict(getixsv(rawcode), tensors)
+    code = optimize_code(rawcode, size_dict, optimizer, simplifier)
+    @show code |> typeof
     TensorNetworkModeling(collect(LT, vars), code, tensors, fixedvertices)
 end
 
@@ -150,7 +152,6 @@ function probability(tn::TensorNetworkModeling; usecuda=false)::AbstractArray
     return tn.code(generate_tensors(tn; usecuda)...)
 end
 
-function OMEinsum.timespacereadwrite_complexity(tn::TensorNetworkModeling)
-    return timespacereadwrite_complexity(tn.code, Dict(zip(get_vars(tn), get_cards(tn; fixedisone=true))))
-end
-OMEinsum.timespace_complexity(tn::TensorNetworkModeling) = timespacereadwrite_complexity(tn)[1:2]
\ No newline at end of file
+function OMEinsum.contraction_complexity(tn::TensorNetworkModeling)
+    return contraction_complexity(tn.code, Dict(zip(get_vars(tn), get_cards(tn; fixedisone=true))))
+end
\ No newline at end of file
diff --git a/src/TensorInference.jl b/src/TensorInference.jl
index 95d9543..8487a3f 100644
--- a/src/TensorInference.jl
+++ b/src/TensorInference.jl
@@ -11,7 +11,7 @@ export timespace_complexity, timespacereadwrite_complexity, TreeSA, GreedyMethod
 export read_uai_file, read_td_file, read_uai_evid_file, read_uai_mar_file, read_uai_problem
 
 # marginals
-export TensorNetworkModeling, get_vars, get_cards, probability, marginals, autorescale!
+export TensorNetworkModeling, get_vars, get_cards, probability, marginals
 
 # MAP
 export most_probable_config, maximum_logp
diff --git a/src/inference.jl b/src/inference.jl
index e0d5031..a3e3700 100644
--- a/src/inference.jl
+++ b/src/inference.jl
@@ -129,16 +129,6 @@ function marginals(tn::TensorNetworkModeling; usecuda=false)::Vector
     vars = get_vars(tn)
     # sometimes, the cost can overflow.
     cost, grads = cost_and_gradient(tn.code, generate_tensors(tn; usecuda))
+    @info "cost = $cost"
     return LinearAlgebra.normalize!.(grads[1:length(vars)], 1)
-end
-
-"""
-$(TYPEDSIGNATURES)
-"""
-function autorescale!(tn::TensorNetworkModeling; usecuda::Bool=false)::TensorNetworkModeling
-    vars = get_vars(tn)
-    for t in tn.tensors[length(vars)+1:end]
-        t ./= maximum(t)/sqrt(length(t))
-    end
-    return tn
 end
\ No newline at end of file
diff --git a/src/utils.jl b/src/utils.jl
index 70f9621..6f484f9 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -76,15 +76,13 @@ function read_uai_evid_file(uai_evid_filepath::AbstractString)
     obsvars = Int64[]
     obsvals = Int64[]
   else
-    # Read the uai evid file into an array of lines
+    # Read the last line of the uai evid file
     line = open(uai_evid_filepath) do file
       readlines(file)
-    end
-
-    @assert length(line) == 1
+    end |> last
 
     # Extract number of observed vars, and their id together with their corresponding value
-    nobsvars, rest = split(line[1]) |> x -> parse.(Int, x) |> x -> (x[1], x[2:end])
+    nobsvars, rest = split(line) |> x -> parse.(Int, x) |> x -> (x[1], x[2:end])
     observations = reshape(rest, 2, :)
 
     # Convert to 1-based indexing
diff --git a/test/inference.jl b/test/inference.jl
index 48fe19b..5c236d6 100644
--- a/test/inference.jl
+++ b/test/inference.jl
@@ -1,5 +1,6 @@
 using Test, Artifacts
 using OMEinsum
+using KaHyPar
 using TensorInference
 
 @testset "gradient-based tensor network solvers" begin
@@ -8,20 +9,20 @@ using TensorInference
 
     # TODO: rescale while contract.
     benchmarks = [
-                  ("Alchemy", 1/3), # fails
-                  ("CSP",1.0),
-                  #("DBN",1.0),
-                  ("Grids",1/2),
-                  ("linkage",1/2),
-                  ("ObjectDetection",1/2),
-                  #("Pedigree",1.0),
-                  ("Promedus",1.0),
-                  #("relational",1.0),
-                  ("Segmentation",1.0),
+                  ("Alchemy", [1/3], TreeSA(ntrials=1, niters=5, βs=0.1:0.1:100)),
+                  ("CSP",fill(1.0, 3), TreeSA(ntrials=1, niters=5, βs=0.1:0.1:100)),
+                  ("DBN",fill(1.0, 6), KaHyParBipartite(sc_target=25)),
+                  ("Grids", [1/4, 1/10, 1/10, 1/32, 1/2, 1/4, 1/15, 1/44], TreeSA(ntrials=1, niters=5, βs=0.1:0.1:100)), # greedy also works
+                  ("linkage", fill(1.2, 17), TreeSA(ntrials=3, niters=20, βs=0.1:0.1:40)),  # linkage_15 fails
+                  ("ObjectDetection",fill(1.5, 65), TreeSA(ntrials=1, niters=5, βs=1:0.1:100)),  # ObjectDetection_35 fails
+                  ("Pedigree",fill(1.0, 3), TreeSA(ntrials=1, niters=5, βs=0.1:0.1:100)), # greedy also works
+                  ("Promedus",fill(1.0, 28), TreeSA(ntrials=1, niters=5, βs=0.1:0.1:100)),  # greedy also works
+                  #("relational",fill(1.0, 5), TreeSA(ntrials=1, niters=5, βs=0.1:0.1:100)),
+                  ("Segmentation", fill(1.0, 6), TreeSA(ntrials=1, niters=5, βs=0.1:0.1:100)),  # greedy also works
                  ]
-    benchmarks = [("linkage",1/3)]
-
-    for (benchmark, rescale) in benchmarks
+    #benchmarks = [("relational", fill(1.0, 5), TreeSA(ntrials=1, niters=5, βs=0.1:0.1:100))]
+    #benchmarks = [("DBN",fill(1.0, 6), SABipartite(sc_target=25, βs=0.1:0.01:50))]
+    for (benchmark, rescales, optimizer) in benchmarks
 
       @testset "$(benchmark) benchmark" begin
 
@@ -31,23 +32,29 @@ using TensorInference
           x -> map(y -> match(rexp, y), x) |> # apply regex
           x -> filter(!isnothing, x) |> # filter out `nothing` values
           x -> map(first, x) # get the first capture of each element
+        @assert length(problems) == length(rescales)
 
-        for problem in problems
+        for (problem, rescale) in [zip(problems, rescales)...]
+          @show problem
 
           @testset "$(problem)" begin
 
             problem = read_uai_problem(problem)
 
             # does not optimize over open vertices
-            tn = TensorNetworkModeling(problem; optimizer=TreeSA(ntrials=3, niters=5, βs=1:0.1:40))
-            autorescale!(tn)
+            #tn = TensorNetworkModeling(problem; optimizer=, rescale)
+            tn = TensorNetworkModeling(problem; optimizer, rescale)
+            sc = contraction_complexity(tn).sc
+            if sc > 28
+              error("space complexity too large! got $(sc)")
+            end
             # @info(tn)
             # @info timespace_complexity(tn)
             marginals2 = marginals(tn)
             # for dangling vertices, the output size is 1.
             npass = 0
             for i=1:problem.nvars
-                npass += (length(marginals2[i]) == 1 && problem.reference_marginals[i] == [0.0, 1]) || isapprox(marginals2[i], problem.reference_marginals[i]; atol=1e-6)
+                npass += length(marginals2[i]) == 1 || isapprox(marginals2[i], problem.reference_marginals[i]; atol=1e-6)
             end
             @test npass == problem.nvars
 

From c689fe08ece2c26eaf16f3d38530c1f41498f03b Mon Sep 17 00:00:00 2001
From: GiggleLiu <cacate0129@gmail.com>
Date: Sat, 8 Oct 2022 00:27:33 +0800
Subject: [PATCH 3/6] clean up a print

---
 src/Core.jl | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/Core.jl b/src/Core.jl
index ca3341c..dcf9271 100644
--- a/src/Core.jl
+++ b/src/Core.jl
@@ -109,7 +109,6 @@ function TensorNetworkModeling(vars::AbstractVector{LT}, rawcode::EinCode, tenso
     # The 3rd and 4th arguments are the optimizer and simplifier that configures which algorithm to use and simplify.
     size_dict = OMEinsum.get_size_dict(getixsv(rawcode), tensors)
     code = optimize_code(rawcode, size_dict, optimizer, simplifier)
-    @show code |> typeof
     TensorNetworkModeling(collect(LT, vars), code, tensors, fixedvertices)
 end
 

From ed11c2754dfbc9a3df0c3a4ead254b58592140b3 Mon Sep 17 00:00:00 2001
From: GiggleLiu <cacate0129@gmail.com>
Date: Sun, 23 Oct 2022 18:20:30 +0800
Subject: [PATCH 4/6] rescaled array

---
 .gitignore             |  1 +
 src/Core.jl            | 25 +++++++------
 src/RescaledArray.jl   | 34 ++++++++++++++++++
 src/TensorInference.jl |  4 ++-
 src/cuda.jl            |  5 ++-
 src/inference.jl       | 29 +++++++++-------
 src/maxprob.jl         |  4 +--
 src/mmap.jl            | 32 +++++++++++------
 test/cuda.jl           | 10 +++---
 test/inference.jl      | 79 ++++++++++++++++++++++++------------------
 test/maxprob.jl        |  4 +--
 test/mmap.jl           |  4 +--
 12 files changed, 151 insertions(+), 80 deletions(-)
 create mode 100644 src/RescaledArray.jl

diff --git a/.gitignore b/.gitignore
index 047c72e..ad6b047 100644
--- a/.gitignore
+++ b/.gitignore
@@ -3,3 +3,4 @@ Manifest.toml
 *.jl.cov
 *.jl.mem
 /docs/build/
+.vscode/
\ No newline at end of file
diff --git a/src/Core.jl b/src/Core.jl
index dcf9271..a7a9717 100644
--- a/src/Core.jl
+++ b/src/Core.jl
@@ -83,20 +83,20 @@ end
 """
 $(TYPEDSIGNATURES)
 """
-function TensorNetworkModeling(instance::UAIInstance; rescale=1.0, openvertices=(), optimizer=GreedyMethod(), simplifier=nothing)::TensorNetworkModeling
-    return TensorNetworkModeling(1:instance.nvars, instance.cards, instance.factors; fixedvertices=Dict(zip(instance.obsvars, instance.obsvals .- 1)), optimizer, simplifier, openvertices, rescale)
+function TensorNetworkModeling(instance::UAIInstance; openvertices=(), optimizer=GreedyMethod(), simplifier=nothing)::TensorNetworkModeling
+    return TensorNetworkModeling(1:instance.nvars, instance.cards, instance.factors; fixedvertices=Dict(zip(instance.obsvars, instance.obsvals .- 1)), optimizer, simplifier, openvertices)
 end
 
 """
 $(TYPEDSIGNATURES)
 """
-function TensorNetworkModeling(vars::AbstractVector{LT}, cards::AbstractVector{Int}, factors::Vector{<:Factor{T}}; rescale=1.0, openvertices=(), fixedvertices=Dict{LT,Int}(), optimizer=GreedyMethod(), simplifier=nothing)::TensorNetworkModeling where {T,LT}
+function TensorNetworkModeling(vars::AbstractVector{LT}, cards::AbstractVector{Int}, factors::Vector{<:Factor{T}}; openvertices=(), fixedvertices=Dict{LT,Int}(), optimizer=GreedyMethod(), simplifier=nothing)::TensorNetworkModeling where {T,LT}
     # The 1st argument of `EinCode` is a vector of vector of labels for specifying the input tensors, 
     # The 2nd argument of `EinCode` is a vector of labels for specifying the output tensor,
     # e.g.
     # `EinCode([[1, 2], [2, 3]], [1, 3])` is the EinCode for matrix multiplication.
     rawcode = EinCode([[[var] for var in vars]..., [[factor.vars...] for factor in factors]...], collect(LT, openvertices))  # labels for vertex tensors (unity tensors) and edge tensors
-    tensors = Array{T}[[ones(T, cards[i]) for i=1:length(vars)]..., [t.vals .* rescale for t in factors]...]
+    tensors = Array{T}[[ones(T, cards[i]) for i=1:length(vars)]..., [t.vals for t in factors]...]
     return TensorNetworkModeling(collect(LT, vars), rawcode, tensors; fixedvertices, optimizer, simplifier)
 end
 """
@@ -134,11 +134,11 @@ chfixedvertices(tn::TensorNetworkModeling, fixedvertices) = TensorNetworkModelin
 """
 $(TYPEDSIGNATURES)
 
-Evaluate the probability of `config`.
+Evaluate the log probability of `config`.
 """
-function probability(tn::TensorNetworkModeling, config)::Real
-    assign = Dict(zip(get_vars(tn), config .+ 1))
-    return mapreduce(x->x[2][getindex.(Ref(assign), x[1])...], *, zip(getixsv(tn.code), tn.tensors))
+function log_probability(tn::TensorNetworkModeling, config::Union{Dict, AbstractVector})::Real
+    assign =  config isa AbstractVector ? Dict(zip(get_vars(tn), config)) : config
+    return sum(x->log(x[2][(getindex.(Ref(assign), x[1]) .+ 1)...]), zip(getixsv(tn.code), tn.tensors))
 end
 
 """
@@ -147,10 +147,13 @@ $(TYPEDSIGNATURES)
 Contract the tensor network and return a probability array with its rank specified in the contraction code `tn.code`.
 The returned array may not be l1-normalized even if the total probability is l1-normalized, because the evidence `tn.fixedvertices` may not be empty.
 """
-function probability(tn::TensorNetworkModeling; usecuda=false)::AbstractArray
-    return tn.code(generate_tensors(tn; usecuda)...)
+function probability(tn::TensorNetworkModeling; usecuda=false, rescale=true)::AbstractArray
+    return tn.code(generate_tensors(tn; usecuda, rescale)...)
 end
 
 function OMEinsum.contraction_complexity(tn::TensorNetworkModeling)
     return contraction_complexity(tn.code, Dict(zip(get_vars(tn), get_cards(tn; fixedisone=true))))
-end
\ No newline at end of file
+end
+
+# adapt array type with the target array type
+match_arraytype(::Type{<:Array{T, N}}, target::AbstractArray{T, N}) where {T, N} = Array(target)
\ No newline at end of file
diff --git a/src/RescaledArray.jl b/src/RescaledArray.jl
new file mode 100644
index 0000000..e6cee4d
--- /dev/null
+++ b/src/RescaledArray.jl
@@ -0,0 +1,34 @@
+# rescaled array: exp(logf) * value
+struct RescaledArray{T, N, AT<:AbstractArray{T, N}} <: AbstractArray{T, N}
+    logf::T
+    value::AT
+end
+Base.show(io::IO, c::RescaledArray) = print(io, "exp($(c.logf)) * $(c.value)")
+Base.show(io::IO, ::MIME"text/plain", c::RescaledArray) = Base.show(io, c)
+Base.Array(c::RescaledArray) = rmul!(Array(c.value), exp(c.logf))
+Base.copy(c::RescaledArray) = RescaledArray(c.logf, copy(c.value))
+
+function rescale_array(tensor::AbstractArray{T}) where T
+    maxf = maximum(tensor)
+    if iszero(maxf)
+        @warn("The maximum value of the array to rescale is 0!")
+        return RescaledArray(zero(T), tensor)
+    end
+    return RescaledArray(log(maxf), OMEinsum.asarray(tensor ./ maxf, tensor))
+end
+function rescale_array(tensor::RescaledArray)
+    res = rescale_array(tensor.value)
+    return RescaledArray(res.logf + tensor.logf, res.value)
+end
+
+for CT in [:DynamicEinCode, :StaticEinCode]
+    @eval function OMEinsum.einsum(code::$CT, @nospecialize(xs::NTuple{N,RescaledArray}), size_dict::Dict) where N
+        res = einsum(code, getfield.(xs, :value), size_dict)
+        return rescale_array(RescaledArray(sum(x->x.logf, xs), res))
+    end
+end
+
+Base.size(arr::RescaledArray) = size(arr.value)
+Base.size(arr::RescaledArray, i::Int) = size(arr.value, i)
+
+match_arraytype(::Type{<:RescaledArray{T, N}}, target::AbstractArray{T, N}) where {T, N} = rescale_array(target)
\ No newline at end of file
diff --git a/src/TensorInference.jl b/src/TensorInference.jl
index 8487a3f..1834ef4 100644
--- a/src/TensorInference.jl
+++ b/src/TensorInference.jl
@@ -5,13 +5,14 @@ using DocStringExtensions, TropicalNumbers
 using Artifacts
 
 # reexport OMEinsum functions
+export RescaledArray
 export timespace_complexity, timespacereadwrite_complexity, TreeSA, GreedyMethod, KaHyParBipartite, SABipartite, MergeGreedy, MergeVectors
 
 # read and load uai files
 export read_uai_file, read_td_file, read_uai_evid_file, read_uai_mar_file, read_uai_problem
 
 # marginals
-export TensorNetworkModeling, get_vars, get_cards, probability, marginals
+export TensorNetworkModeling, get_vars, get_cards, log_probability, probability, marginals
 
 # MAP
 export most_probable_config, maximum_logp
@@ -20,6 +21,7 @@ export most_probable_config, maximum_logp
 export MMAPModeling
 
 include("Core.jl")
+include("RescaledArray.jl")
 include("utils.jl")
 include("inference.jl")
 include("maxprob.jl")
diff --git a/src/cuda.jl b/src/cuda.jl
index 0e1f3a2..964f9eb 100644
--- a/src/cuda.jl
+++ b/src/cuda.jl
@@ -4,4 +4,7 @@ function onehot_like(A::CuArray, j)
     mask = zero(A)
     CUDA.@allowscalar mask[j] = one(eltype(mask))
     return mask
-end
\ No newline at end of file
+end
+
+# NOTE: this interface should be in OMEinsum
+match_arraytype(::Type{<:CuArray{T, N}}, target::AbstractArray{T, N}) where {T, N} = CuArray(target)
\ No newline at end of file
diff --git a/src/inference.jl b/src/inference.jl
index a3e3700..45b8ab8 100644
--- a/src/inference.jl
+++ b/src/inference.jl
@@ -1,21 +1,22 @@
 # generate tensors based on which vertices are fixed.
-generate_tensors(gp::TensorNetworkModeling; usecuda) = generate_tensors(gp.code, gp.tensors, gp.fixedvertices; usecuda)
-function generate_tensors(code, tensors, fixedvertices; usecuda)
-    isempty(fixedvertices) && return tensors
+generate_tensors(gp::TensorNetworkModeling; usecuda, rescale) = generate_tensors(gp.code, gp.tensors, gp.fixedvertices; usecuda, rescale)
+function generate_tensors(code, tensors, fixedvertices; usecuda, rescale)
     ixs = getixsv(code)
     # `ix` is the vector of labels (or a degree of freedoms) for a tensor,
     # if a label in `ix` is fixed to a value, do the slicing to the tensor it associates to.
     map(tensors, ixs) do t, ix
         dims = map(ixi->ixi ∉ keys(fixedvertices) ? Colon() : (fixedvertices[ixi]+1:fixedvertices[ixi]+1), ix)
-        usecuda ? CuArray(t[dims...]) : t[dims...]
+        t2 = t[dims...]
+        t3 = usecuda ? CuArray(t2) : t2
+        rescale ? rescale_array(t3) : t3
     end
 end
 
 # ######### Inference by back propagation ############
 # `CacheTree` stores intermediate `NestedEinsum` contraction results.
 # It is a tree structure that isomorphic to the contraction tree,
-# `siblings` are the siblings of current node.
 # `content` is the cached intermediate contraction result.
+# `siblings` are the siblings of current node.
 struct CacheTree{T}
     content::AbstractArray{T}
     siblings::Vector{CacheTree{T}}
@@ -60,7 +61,7 @@ function generate_gradient_tree(code::NestedEinsum, cache::CacheTree{T}, dy::Abs
     if OMEinsum.isleaf(code)
         return CacheTree(dy, CacheTree{T}[])
     else
-        xs = (getfield.(cache.siblings, :content)...,)
+        xs = ntuple(i->cache.siblings[i].content, length(cache.siblings))
         # `einsum_grad` is the back-propagation rule for einsum function.
         # If the forward pass is `y = einsum(EinCode(inputs_labels, output_labels), (A, B, ...), size_dict)`
         # Then the back-propagation pass is
@@ -87,7 +88,7 @@ function gradient_tree(code, xs)
     # forward compute and cache intermediate results.
     cache = cached_einsum(code, xs, size_dict)
     # initialize `y̅` as `1`. Note we always start from `L̅ := 1`.
-    dy = fill!(similar(cache.content), one(eltype(cache.content)))
+    dy = match_arraytype(typeof(cache.content), ones(eltype(cache.content), size(cache.content)))
     # back-propagate
     return copy(cache.content), generate_gradient_tree(code, cache, dy, size_dict)
 end
@@ -125,10 +126,14 @@ $(TYPEDSIGNATURES)
 Returns the marginal probability distribution of variables.
 One can use `get_vars(tn)` to get the full list of variables in this tensor network.
 """
-function marginals(tn::TensorNetworkModeling; usecuda=false)::Vector
+function marginals(tn::TensorNetworkModeling; usecuda=false, rescale=true)::Vector
     vars = get_vars(tn)
-    # sometimes, the cost can overflow.
-    cost, grads = cost_and_gradient(tn.code, generate_tensors(tn; usecuda))
-    @info "cost = $cost"
-    return LinearAlgebra.normalize!.(grads[1:length(vars)], 1)
+    # sometimes, the cost can overflow, then we need to rescale the tensors during contraction.
+    cost, grads = cost_and_gradient(tn.code, generate_tensors(tn; usecuda, rescale))
+    @debug "cost = $cost"
+    if rescale
+        return LinearAlgebra.normalize!.(getfield.(grads[1:length(vars)], :value), 1)
+    else
+        return LinearAlgebra.normalize!.(grads[1:length(vars)], 1)
+    end
 end
\ No newline at end of file
diff --git a/src/maxprob.jl b/src/maxprob.jl
index 7117a4c..ea16581 100644
--- a/src/maxprob.jl
+++ b/src/maxprob.jl
@@ -46,7 +46,7 @@ Returns the largest log-probability and the most probable configuration.
 """
 function most_probable_config(tn::TensorNetworkModeling; usecuda=false)::Tuple{Tropical,Vector}
     vars = get_vars(tn)
-    tensors = map(t->Tropical.(log.(t)), generate_tensors(tn; usecuda))
+    tensors = map(t->Tropical.(log.(t)), generate_tensors(tn; usecuda, rescale=false))
     logp, grads = cost_and_gradient(tn.code, tensors)
     # use Array to convert CuArray to CPU arrays
     return Array(logp)[], map(k->haskey(tn.fixedvertices, vars[k]) ? tn.fixedvertices[vars[k]] : argmax(grads[k]) - 1, 1:length(vars))
@@ -59,6 +59,6 @@ Returns an output array containing largest log-probabilities.
 """
 function maximum_logp(tn::TensorNetworkModeling; usecuda=false)::AbstractArray{<:Tropical}
     # generate tropical tensors with its elements being log(p).
-    tensors = map(t->Tropical.(log.(t)), generate_tensors(tn; usecuda))
+    tensors = map(t->Tropical.(log.(t)), generate_tensors(tn; usecuda, rescale=false))
     return tn.code(tensors...)
 end
\ No newline at end of file
diff --git a/src/mmap.jl b/src/mmap.jl
index 1ff2264..5ccb5e0 100644
--- a/src/mmap.jl
+++ b/src/mmap.jl
@@ -122,7 +122,10 @@ function OMEinsum.timespacereadwrite_complexity(mmap::MMAPModeling{LT}) where LT
 end
 OMEinsum.timespace_complexity(mmap::MMAPModeling) = timespacereadwrite_complexity(mmap)[1:2]
 
-generate_tensors(mmap::MMAPModeling; usecuda) = [generate_tensors(mmap.code, mmap.tensors, mmap.fixedvertices; usecuda)..., map(cluster->probability(cluster; fixedvertices=mmap.fixedvertices, usecuda), mmap.clusters)...]
+function generate_tensors(mmap::MMAPModeling; usecuda, rescale)
+    return [generate_tensors(mmap.code, mmap.tensors, mmap.fixedvertices; usecuda, rescale)...,
+        map(cluster->probability(cluster; fixedvertices=mmap.fixedvertices, usecuda, rescale), mmap.clusters)...]
+end
 
 # find connected clusters
 function connected_clusters(ixs, vars::Vector{LT}) where LT
@@ -174,7 +177,7 @@ $(TYPEDSIGNATURES)
 """
 function most_probable_config(mmap::MMAPModeling; usecuda=false)::Tuple{Tropical,Vector}
     vars = get_vars(mmap)
-    tensors = map(t->OMEinsum.asarray(Tropical.(log.(t)), t), generate_tensors(mmap; usecuda))
+    tensors = map(t->OMEinsum.asarray(Tropical.(log.(t)), t), generate_tensors(mmap; usecuda, rescale=false))
     logp, grads = cost_and_gradient(mmap.code, tensors)
     # use Array to convert CuArray to CPU arrays
     return Array(logp)[], map(k->haskey(mmap.fixedvertices, vars[k]) ? mmap.fixedvertices[vars[k]] : argmax(grads[k]) - 1, 1:length(vars))
@@ -184,22 +187,29 @@ end
 $(TYPEDSIGNATURES)
 """
 function maximum_logp(mmap::MMAPModeling; usecuda=false)::AbstractArray{<:Tropical}
-    tensors = map(t->OMEinsum.asarray(Tropical.(log.(t)), t), generate_tensors(mmap; usecuda))
+    tensors = map(t->OMEinsum.asarray(Tropical.(log.(t)), t), generate_tensors(mmap; usecuda, rescale=false))
     return mmap.code(tensors...)
 end
 
 """
 $(TYPEDSIGNATURES)
 """
-function probability(mmap::MMAPModeling, config)::Real
-    fixedvertices = merge(mmap.fixedvertices, Dict(zip(get_vars(mmap), config)))
-    assign = Dict(zip(get_vars(mmap), config .+ 1))
-    m1 = mapreduce(x->x[2][getindex.(Ref(assign), x[1])...], *, zip(getixsv(mmap.code), mmap.tensors))
-    m2 = prod(cluster->probability(cluster; fixedvertices, usecuda=false)[], mmap.clusters)
-    return m1 * m2
+function log_probability(mmap::MMAPModeling, config::Union{Dict, AbstractVector}; rescale=true, usecuda=false)::Real
+    @assert length(get_vars(mmap)) == length(config)
+    fixedvertices =  config isa AbstractVector ? Dict(zip(get_vars(mmap), config)) : config
+    assign = merge(mmap.fixedvertices, fixedvertices)
+    # two contributions to the probability, not-clustered tensors and clusters.
+    m1 = sum(x->log(x[2][(getindex.(Ref(assign), x[1]) .+ 1)...]), zip(getixsv(mmap.code), mmap.tensors))
+    m2 = sum(cluster->probability(cluster; fixedvertices, usecuda, rescale).logf, mmap.clusters)
+    return m1 + m2
 end
 
-function probability(c::Cluster; fixedvertices, usecuda)::AbstractArray
-    tensors = generate_tensors(c.code, c.tensors, fixedvertices; usecuda)
+function probability(c::Cluster; fixedvertices, usecuda, rescale)::AbstractArray
+    tensors = generate_tensors(c.code, c.tensors, fixedvertices; usecuda, rescale)
     return c.code(tensors...)
+end
+
+function log_probability(c::Cluster, config::Union{AbstractVector, Dict})::AbstractArray
+    assign =  config isa AbstractVector ? Dict(zip(get_vars(c), config)) : config
+    return sum(x->log(x[2][(getindex.(Ref(assign), x[1]) .+ 1)...]), zip(getixsv(c.code), c.tensors))
 end
\ No newline at end of file
diff --git a/test/cuda.jl b/test/cuda.jl
index 1367e62..8c0e38e 100644
--- a/test/cuda.jl
+++ b/test/cuda.jl
@@ -9,7 +9,7 @@ CUDA.allowscalar(false)
 
     # does not optimize over open vertices
     tn = TensorNetworkModeling(instance; optimizer=TreeSA(ntrials=1, niters=2, βs=1:0.1:40))
-    @info timespace_complexity(tn)
+    @info contraction_complexity(tn)
     @time marginals2 = marginals(tn; usecuda=true)
     @test all(x->x isa CuArray, marginals2)
     # for dangling vertices, the output size is 1.
@@ -26,10 +26,10 @@ end
 
     # does not optimize over open vertices
     tn = TensorNetworkModeling(instance; optimizer=TreeSA(ntrials=1, niters=2, βs=1:0.1:40))
-    @info timespace_complexity(tn)
+    @info contraction_complexity(tn)
     most_probable_config(tn)
     @time logp, config = most_probable_config(tn; usecuda=true)
-    @test probability(tn, config) ≈ exp(logp.n)
+    @test log_probability(tn, config) ≈ logp.n
     culogp = maximum_logp(tn; usecuda=true)
     @test culogp isa CuArray
     @test Array(culogp)[] ≈ logp
@@ -52,12 +52,12 @@ end
     tn2 = MMAPModeling(instance; marginalizedvertices=collect(1:instance.nvars), optimizer)
     cup = probability(tn_ref; usecuda=true)
     culogp = maximum_logp(tn2; usecuda=true)
-    @test cup isa CuArray
+    @test cup isa RescaledArray{T, N, <:CuArray} where {T, N}
     @test culogp isa CuArray
     @test Array(cup)[] ≈ exp(Array(culogp)[].n)
 
     # does not optimize over open vertices
     tn3 = MMAPModeling(instance; marginalizedvertices=[2,4,6], optimizer)
     logp, config = most_probable_config(tn3; usecuda=true)
-    @test probability(tn3, config) ≈ exp(logp.n)
+    @test log_probability(tn3, config) ≈ logp.n
 end 
diff --git a/test/inference.jl b/test/inference.jl
index 5c236d6..cccb0ff 100644
--- a/test/inference.jl
+++ b/test/inference.jl
@@ -3,26 +3,58 @@ using OMEinsum
 using KaHyPar
 using TensorInference
 
+@testset "composite number" begin
+    A = RescaledArray(2.0, [2.0 3.0; 5.0 6.0])
+    x = RescaledArray(2.0, [2.0, 3.0])
+    op = ein"ij, j -> i"
+    @test Array(x) ≈ exp(2.0) .* [2.0, 3.0]
+    @test op(Array(A), Array(x)) ≈ Array(op(A, x))
+    println(x)
+end
+
+@testset "cached, rescaled contract" begin
+    problem = read_uai_problem("Promedus_14")
+    optimizer = TreeSA(ntrials=1, niters=5, βs=0.1:0.1:100)
+    tn = TensorNetworkModeling(problem; optimizer)
+    p1 = probability(tn; usecuda=false, rescale=false)
+    p2 = probability(tn; usecuda=false, rescale=true)
+    @test p1 ≈ Array(p2)
+    
+    # cached contract
+    xs = TensorInference.generate_tensors(tn; usecuda=false, rescale=true)
+    size_dict = OMEinsum.get_size_dict!(getixsv(tn.code), xs, Dict{Int,Int}())
+    cache = TensorInference.cached_einsum(tn.code, xs, size_dict)
+    @test cache.content isa RescaledArray
+    @test Array(cache.content) ≈ p1
+
+    # compute marginals
+    marginals2 = marginals(tn)
+    npass = 0
+    for i=1:problem.nvars
+        npass += length(marginals2[i]) == 1 || isapprox(marginals2[i], problem.reference_marginals[i]; atol=1e-6)
+    end
+    @test npass == problem.nvars
+end
+
 @testset "gradient-based tensor network solvers" begin
 
   @testset "UAI 2014 problem set" begin
 
-    # TODO: rescale while contract.
     benchmarks = [
-                  ("Alchemy", [1/3], TreeSA(ntrials=1, niters=5, βs=0.1:0.1:100)),
-                  ("CSP",fill(1.0, 3), TreeSA(ntrials=1, niters=5, βs=0.1:0.1:100)),
-                  ("DBN",fill(1.0, 6), KaHyParBipartite(sc_target=25)),
-                  ("Grids", [1/4, 1/10, 1/10, 1/32, 1/2, 1/4, 1/15, 1/44], TreeSA(ntrials=1, niters=5, βs=0.1:0.1:100)), # greedy also works
-                  ("linkage", fill(1.2, 17), TreeSA(ntrials=3, niters=20, βs=0.1:0.1:40)),  # linkage_15 fails
-                  ("ObjectDetection",fill(1.5, 65), TreeSA(ntrials=1, niters=5, βs=1:0.1:100)),  # ObjectDetection_35 fails
-                  ("Pedigree",fill(1.0, 3), TreeSA(ntrials=1, niters=5, βs=0.1:0.1:100)), # greedy also works
-                  ("Promedus",fill(1.0, 28), TreeSA(ntrials=1, niters=5, βs=0.1:0.1:100)),  # greedy also works
-                  #("relational",fill(1.0, 5), TreeSA(ntrials=1, niters=5, βs=0.1:0.1:100)),
-                  ("Segmentation", fill(1.0, 6), TreeSA(ntrials=1, niters=5, βs=0.1:0.1:100)),  # greedy also works
+                  #("Alchemy", TreeSA(ntrials=1, niters=5, βs=0.1:0.1:100)),
+                  #("CSP", TreeSA(ntrials=1, niters=5, βs=0.1:0.1:100)),
+                  #("DBN", KaHyParBipartite(sc_target=25)),
+                  #("Grids", TreeSA(ntrials=1, niters=5, βs=0.1:0.1:100)), # greedy also works
+                  #("linkage", TreeSA(ntrials=3, niters=20, βs=0.1:0.1:40)),  # linkage_15 fails
+                  #("ObjectDetection", TreeSA(ntrials=1, niters=5, βs=1:0.1:100)),  # ObjectDetection_35 fails
+                  #("Pedigree", TreeSA(ntrials=1, niters=5, βs=0.1:0.1:100)), # greedy also works
+                  ("Promedus", TreeSA(ntrials=1, niters=5, βs=0.1:0.1:100)),  # greedy also works
+                  #("relational", TreeSA(ntrials=1, niters=5, βs=0.1:0.1:100)),
+                  ("Segmentation", TreeSA(ntrials=1, niters=5, βs=0.1:0.1:100)),  # greedy also works
                  ]
     #benchmarks = [("relational", fill(1.0, 5), TreeSA(ntrials=1, niters=5, βs=0.1:0.1:100))]
     #benchmarks = [("DBN",fill(1.0, 6), SABipartite(sc_target=25, βs=0.1:0.01:50))]
-    for (benchmark, rescales, optimizer) in benchmarks
+    for (benchmark, optimizer) in benchmarks
 
       @testset "$(benchmark) benchmark" begin
 
@@ -32,9 +64,8 @@ using TensorInference
           x -> map(y -> match(rexp, y), x) |> # apply regex
           x -> filter(!isnothing, x) |> # filter out `nothing` values
           x -> map(first, x) # get the first capture of each element
-        @assert length(problems) == length(rescales)
 
-        for (problem, rescale) in [zip(problems, rescales)...]
+        for problem in problems
           @show problem
 
           @testset "$(problem)" begin
@@ -42,8 +73,7 @@ using TensorInference
             problem = read_uai_problem(problem)
 
             # does not optimize over open vertices
-            #tn = TensorNetworkModeling(problem; optimizer=, rescale)
-            tn = TensorNetworkModeling(problem; optimizer, rescale)
+            tn = TensorNetworkModeling(problem; optimizer)
             sc = contraction_complexity(tn).sc
             if sc > 28
               error("space complexity too large! got $(sc)")
@@ -63,21 +93,4 @@ using TensorInference
       end
     end
   end
-
-  # problem_name = "Promedus_14"
-  # @testset "$problem_name" begin
-  #   problem = read_uai_problem(problem_name)
-  #   # does not optimize over open vertices
-  #   tn = TensorNetworkModeling(problem; optimizer=TreeSA(ntrials=1, niters=2, βs=1:0.1:40))
-  #   # @info(tn)
-  #   # @info timespace_complexity(tn)
-  #   marginals2 = marginals(tn)
-  #   # for dangling vertices, the output size is 1.
-  #   npass = 0
-  #   for i=1:problem.nvars
-  #       npass += (length(marginals2[i]) == 1 && problem.reference_marginals[i] == [0.0, 1]) || isapprox(marginals2[i], problem.reference_marginals[i]; atol=1e-6)
-  #   end
-  #   @test npass == problem.nvars
-  # end
-
 end 
diff --git a/test/maxprob.jl b/test/maxprob.jl
index 9bdfe5f..0f5fafd 100644
--- a/test/maxprob.jl
+++ b/test/maxprob.jl
@@ -8,9 +8,9 @@ using TensorInference
 
     # does not optimize over open vertices
     tn = TensorNetworkModeling(instance; optimizer=TreeSA(ntrials=1, niters=2, βs=1:0.1:40))
-    @info timespace_complexity(tn)
+    @info contraction_complexity(tn)
     most_probable_config(tn)
     @time logp, config = most_probable_config(tn)
-    @test probability(tn, config) ≈ exp(logp.n)
+    @test log_probability(tn, config) ≈ logp.n
     @test maximum_logp(tn)[] ≈ logp
 end 
diff --git a/test/mmap.jl b/test/mmap.jl
index c6fad7a..5a38590 100644
--- a/test/mmap.jl
+++ b/test/mmap.jl
@@ -21,11 +21,11 @@ end
     # marginalize all vars
     mmap2 = MMAPModeling(instance; marginalizedvertices=collect(1:instance.nvars), optimizer)
     @info(mmap2)
-    @test probability(tn_ref)[] ≈ exp(maximum_logp(mmap2)[].n)
+    @test Array(probability(tn_ref))[] ≈ exp(maximum_logp(mmap2)[].n)
 
     # does not optimize over open vertices
     mmap3 = MMAPModeling(instance; marginalizedvertices=[2,4,6], optimizer)
     @info(mmap3)
     logp, config = most_probable_config(mmap3)
-    @test probability(mmap3, config) ≈ exp(logp.n)
+    @test log_probability(mmap3, config) ≈ logp.n
 end 

From 863d44d3360f10a226540898a2c2b88fef8dc81f Mon Sep 17 00:00:00 2001
From: GiggleLiu <cacate0129@gmail.com>
Date: Thu, 3 Nov 2022 21:37:37 +0800
Subject: [PATCH 5/6] update

---
 src/Core.jl            | 32 +++++++++++++++---------------
 src/RescaledArray.jl   | 43 +++++++++++++++++++++++++++--------------
 src/TensorInference.jl |  4 ++--
 src/inference.jl       | 10 +++++-----
 src/maxprob.jl         |  8 ++++----
 src/mmap.jl            | 44 +++++++++++++++++++++---------------------
 test/cuda.jl           | 12 ++++++------
 test/inference.jl      |  6 +++---
 test/maxprob.jl        |  2 +-
 test/mmap.jl           |  8 ++++----
 10 files changed, 91 insertions(+), 78 deletions(-)

diff --git a/src/Core.jl b/src/Core.jl
index a7a9717..c1581fc 100644
--- a/src/Core.jl
+++ b/src/Core.jl
@@ -48,14 +48,14 @@ Probabilistic modeling with a tensor network.
 * `tensors` is the tensors fed into the tensor network.
 * `fixedvertices` is a dictionary to specifiy degree of freedoms fixed to certain values.
 """
-struct TensorNetworkModeling{LT,ET,MT<:AbstractArray}
+struct TensorNetworkModel{LT,ET,MT<:AbstractArray}
     vars::Vector{LT}
     code::ET
     tensors::Vector{MT}
     fixedvertices::Dict{LT,Int}
 end
 
-function Base.show(io::IO, tn::TensorNetworkModeling)
+function Base.show(io::IO, tn::TensorNetworkModel)
     open = getiyv(tn.code)
     variables = join([string_var(var, open, tn.fixedvertices) for var in tn.vars], ", ")
     tc, sc, rw = timespacereadwrite_complexity(tn)
@@ -63,7 +63,7 @@ function Base.show(io::IO, tn::TensorNetworkModeling)
     println(io, "variables: $variables")
     print_tcscrw(io, tc, sc, rw)
 end
-Base.show(io::IO, ::MIME"text/plain", tn::TensorNetworkModeling) = Base.show(io, tn)
+Base.show(io::IO, ::MIME"text/plain", tn::TensorNetworkModel) = Base.show(io, tn)
 function string_var(var, open ,fixedvertices)
     if var ∈ open && haskey(fixedvertices, var)
         "$var (open, fixed to $(fixedvertices[var]))"
@@ -83,33 +83,33 @@ end
 """
 $(TYPEDSIGNATURES)
 """
-function TensorNetworkModeling(instance::UAIInstance; openvertices=(), optimizer=GreedyMethod(), simplifier=nothing)::TensorNetworkModeling
-    return TensorNetworkModeling(1:instance.nvars, instance.cards, instance.factors; fixedvertices=Dict(zip(instance.obsvars, instance.obsvals .- 1)), optimizer, simplifier, openvertices)
+function TensorNetworkModel(instance::UAIInstance; openvertices=(), optimizer=GreedyMethod(), simplifier=nothing)::TensorNetworkModel
+    return TensorNetworkModel(1:instance.nvars, instance.cards, instance.factors; fixedvertices=Dict(zip(instance.obsvars, instance.obsvals .- 1)), optimizer, simplifier, openvertices)
 end
 
 """
 $(TYPEDSIGNATURES)
 """
-function TensorNetworkModeling(vars::AbstractVector{LT}, cards::AbstractVector{Int}, factors::Vector{<:Factor{T}}; openvertices=(), fixedvertices=Dict{LT,Int}(), optimizer=GreedyMethod(), simplifier=nothing)::TensorNetworkModeling where {T,LT}
+function TensorNetworkModel(vars::AbstractVector{LT}, cards::AbstractVector{Int}, factors::Vector{<:Factor{T}}; openvertices=(), fixedvertices=Dict{LT,Int}(), optimizer=GreedyMethod(), simplifier=nothing)::TensorNetworkModel where {T,LT}
     # The 1st argument of `EinCode` is a vector of vector of labels for specifying the input tensors, 
     # The 2nd argument of `EinCode` is a vector of labels for specifying the output tensor,
     # e.g.
     # `EinCode([[1, 2], [2, 3]], [1, 3])` is the EinCode for matrix multiplication.
     rawcode = EinCode([[[var] for var in vars]..., [[factor.vars...] for factor in factors]...], collect(LT, openvertices))  # labels for vertex tensors (unity tensors) and edge tensors
     tensors = Array{T}[[ones(T, cards[i]) for i=1:length(vars)]..., [t.vals for t in factors]...]
-    return TensorNetworkModeling(collect(LT, vars), rawcode, tensors; fixedvertices, optimizer, simplifier)
+    return TensorNetworkModel(collect(LT, vars), rawcode, tensors; fixedvertices, optimizer, simplifier)
 end
 """
 $(TYPEDSIGNATURES)
 """
-function TensorNetworkModeling(vars::AbstractVector{LT}, rawcode::EinCode, tensors::Vector{<:AbstractArray}; fixedvertices=Dict{LT,Int}(), optimizer=GreedyMethod(), simplifier=nothing)::TensorNetworkModeling where LT
+function TensorNetworkModel(vars::AbstractVector{LT}, rawcode::EinCode, tensors::Vector{<:AbstractArray}; fixedvertices=Dict{LT,Int}(), optimizer=GreedyMethod(), simplifier=nothing)::TensorNetworkModel where LT
     # `optimize_code` optimizes the contraction order of a raw tensor network without a contraction order specified.
     # The 1st argument is the contraction pattern to be optimized (without contraction order).
     # The 2nd arugment is the size dictionary, which is a label-integer dictionary.
     # The 3rd and 4th arguments are the optimizer and simplifier that configures which algorithm to use and simplify.
     size_dict = OMEinsum.get_size_dict(getixsv(rawcode), tensors)
     code = optimize_code(rawcode, size_dict, optimizer, simplifier)
-    TensorNetworkModeling(collect(LT, vars), code, tensors, fixedvertices)
+    TensorNetworkModel(collect(LT, vars), code, tensors, fixedvertices)
 end
 
 """
@@ -117,26 +117,26 @@ $(TYPEDSIGNATURES)
 
 Get the variables in this tensor network, they are also known as legs, labels, or degree of freedoms.
 """
-get_vars(tn::TensorNetworkModeling)::Vector = tn.vars
+get_vars(tn::TensorNetworkModel)::Vector = tn.vars
 
 """
 $(TYPEDSIGNATURES)
 
 Get the cardinalities of variables in this tensor network.
 """
-function get_cards(tn::TensorNetworkModeling; fixedisone=false)::Vector
+function get_cards(tn::TensorNetworkModel; fixedisone=false)::Vector
     vars = get_vars(tn)
     [fixedisone && haskey(tn.fixedvertices, vars[k]) ? 1 : length(tn.tensors[k]) for k=1:length(vars)]
 end
 
-chfixedvertices(tn::TensorNetworkModeling, fixedvertices) = TensorNetworkModeling(tn.vars, tn.code, tn.tensors, fixedvertices)
+chfixedvertices(tn::TensorNetworkModel, fixedvertices) = TensorNetworkModel(tn.vars, tn.code, tn.tensors, fixedvertices)
 
 """
 $(TYPEDSIGNATURES)
 
 Evaluate the log probability of `config`.
 """
-function log_probability(tn::TensorNetworkModeling, config::Union{Dict, AbstractVector})::Real
+function log_probability(tn::TensorNetworkModel, config::Union{Dict, AbstractVector})::Real
     assign =  config isa AbstractVector ? Dict(zip(get_vars(tn), config)) : config
     return sum(x->log(x[2][(getindex.(Ref(assign), x[1]) .+ 1)...]), zip(getixsv(tn.code), tn.tensors))
 end
@@ -147,11 +147,11 @@ $(TYPEDSIGNATURES)
 Contract the tensor network and return a probability array with its rank specified in the contraction code `tn.code`.
 The returned array may not be l1-normalized even if the total probability is l1-normalized, because the evidence `tn.fixedvertices` may not be empty.
 """
-function probability(tn::TensorNetworkModeling; usecuda=false, rescale=true)::AbstractArray
-    return tn.code(generate_tensors(tn; usecuda, rescale)...)
+function probability(tn::TensorNetworkModel; usecuda=false, rescale=true)::AbstractArray
+    return tn.code(adapt_tensors(tn; usecuda, rescale)...)
 end
 
-function OMEinsum.contraction_complexity(tn::TensorNetworkModeling)
+function OMEinsum.contraction_complexity(tn::TensorNetworkModel)
     return contraction_complexity(tn.code, Dict(zip(get_vars(tn), get_cards(tn; fixedisone=true))))
 end
 
diff --git a/src/RescaledArray.jl b/src/RescaledArray.jl
index e6cee4d..97f66d4 100644
--- a/src/RescaledArray.jl
+++ b/src/RescaledArray.jl
@@ -1,14 +1,26 @@
-# rescaled array: exp(logf) * value
+"""
+$(TYPEDEF)
+    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 RescaledArray{T, N, AT<:AbstractArray{T, N}} <: AbstractArray{T, N}
-    logf::T
-    value::AT
+    log_factor::T
+    normalized_value::AT
 end
-Base.show(io::IO, c::RescaledArray) = print(io, "exp($(c.logf)) * $(c.value)")
+Base.show(io::IO, c::RescaledArray) = print(io, "exp($(c.log_factor)) * $(c.normalized_value)")
 Base.show(io::IO, ::MIME"text/plain", c::RescaledArray) = Base.show(io, c)
-Base.Array(c::RescaledArray) = rmul!(Array(c.value), exp(c.logf))
-Base.copy(c::RescaledArray) = RescaledArray(c.logf, copy(c.value))
+Base.Array(c::RescaledArray) = rmul!(Array(c.normalized_value), exp(c.log_factor))
+Base.copy(c::RescaledArray) = RescaledArray(c.log_factor, copy(c.normalized_value))
+
+"""
+$(TYPEDSIGNATURES)
 
-function rescale_array(tensor::AbstractArray{T}) where T
+Returns a rescaled array that equivalent to the input tensor.
+"""
+function rescale_array(tensor::AbstractArray{T})::RescaledArray where T
     maxf = maximum(tensor)
     if iszero(maxf)
         @warn("The maximum value of the array to rescale is 0!")
@@ -16,19 +28,20 @@ function rescale_array(tensor::AbstractArray{T}) where T
     end
     return RescaledArray(log(maxf), OMEinsum.asarray(tensor ./ maxf, tensor))
 end
-function rescale_array(tensor::RescaledArray)
-    res = rescale_array(tensor.value)
-    return RescaledArray(res.logf + tensor.logf, res.value)
-end
 
 for CT in [:DynamicEinCode, :StaticEinCode]
     @eval function OMEinsum.einsum(code::$CT, @nospecialize(xs::NTuple{N,RescaledArray}), size_dict::Dict) where N
-        res = einsum(code, getfield.(xs, :value), size_dict)
-        return rescale_array(RescaledArray(sum(x->x.logf, xs), res))
+        # The following equality holds
+        # einsum(code, exp(α) * A, exp(β) * B, ...) = exp(α * β * ...) * einsum(code, A, B, ...)
+        # Hence the einsum is performed on the normalized values, and the factors are added later.
+        res = einsum(code, getfield.(xs, :normalized_value), size_dict)
+        rescaled = rescale_array(res)
+        # a new rescaled array, its factor is 
+        return RescaledArray(sum(x->x.log_factor, xs) + rescaled.log_factor, rescaled.normalized_value)
     end
 end
 
-Base.size(arr::RescaledArray) = size(arr.value)
-Base.size(arr::RescaledArray, i::Int) = size(arr.value, i)
+Base.size(arr::RescaledArray) = size(arr.normalized_value)
+Base.size(arr::RescaledArray, i::Int) = size(arr.normalized_value, i)
 
 match_arraytype(::Type{<:RescaledArray{T, N}}, target::AbstractArray{T, N}) where {T, N} = rescale_array(target)
\ No newline at end of file
diff --git a/src/TensorInference.jl b/src/TensorInference.jl
index 1834ef4..94053f1 100644
--- a/src/TensorInference.jl
+++ b/src/TensorInference.jl
@@ -12,13 +12,13 @@ export timespace_complexity, timespacereadwrite_complexity, TreeSA, GreedyMethod
 export read_uai_file, read_td_file, read_uai_evid_file, read_uai_mar_file, read_uai_problem
 
 # marginals
-export TensorNetworkModeling, get_vars, get_cards, log_probability, probability, marginals
+export TensorNetworkModel, get_vars, get_cards, log_probability, probability, marginals
 
 # MAP
 export most_probable_config, maximum_logp
 
 # MMAP
-export MMAPModeling
+export MMAPModel
 
 include("Core.jl")
 include("RescaledArray.jl")
diff --git a/src/inference.jl b/src/inference.jl
index 45b8ab8..804edbf 100644
--- a/src/inference.jl
+++ b/src/inference.jl
@@ -1,6 +1,6 @@
 # generate tensors based on which vertices are fixed.
-generate_tensors(gp::TensorNetworkModeling; usecuda, rescale) = generate_tensors(gp.code, gp.tensors, gp.fixedvertices; usecuda, rescale)
-function generate_tensors(code, tensors, fixedvertices; usecuda, rescale)
+adapt_tensors(gp::TensorNetworkModel; usecuda, rescale) = adapt_tensors(gp.code, gp.tensors, gp.fixedvertices; usecuda, rescale)
+function adapt_tensors(code, tensors, fixedvertices; usecuda, rescale)
     ixs = getixsv(code)
     # `ix` is the vector of labels (or a degree of freedoms) for a tensor,
     # if a label in `ix` is fixed to a value, do the slicing to the tensor it associates to.
@@ -126,13 +126,13 @@ $(TYPEDSIGNATURES)
 Returns the marginal probability distribution of variables.
 One can use `get_vars(tn)` to get the full list of variables in this tensor network.
 """
-function marginals(tn::TensorNetworkModeling; usecuda=false, rescale=true)::Vector
+function marginals(tn::TensorNetworkModel; usecuda=false, rescale=true)::Vector
     vars = get_vars(tn)
     # sometimes, the cost can overflow, then we need to rescale the tensors during contraction.
-    cost, grads = cost_and_gradient(tn.code, generate_tensors(tn; usecuda, rescale))
+    cost, grads = cost_and_gradient(tn.code, adapt_tensors(tn; usecuda, rescale))
     @debug "cost = $cost"
     if rescale
-        return LinearAlgebra.normalize!.(getfield.(grads[1:length(vars)], :value), 1)
+        return LinearAlgebra.normalize!.(getfield.(grads[1:length(vars)], :rescaled_value), 1)
     else
         return LinearAlgebra.normalize!.(grads[1:length(vars)], 1)
     end
diff --git a/src/maxprob.jl b/src/maxprob.jl
index ea16581..e947de3 100644
--- a/src/maxprob.jl
+++ b/src/maxprob.jl
@@ -44,9 +44,9 @@ $(TYPEDSIGNATURES)
 
 Returns the largest log-probability and the most probable configuration.
 """
-function most_probable_config(tn::TensorNetworkModeling; usecuda=false)::Tuple{Tropical,Vector}
+function most_probable_config(tn::TensorNetworkModel; usecuda=false)::Tuple{Tropical,Vector}
     vars = get_vars(tn)
-    tensors = map(t->Tropical.(log.(t)), generate_tensors(tn; usecuda, rescale=false))
+    tensors = map(t->Tropical.(log.(t)), adapt_tensors(tn; usecuda, rescale=false))
     logp, grads = cost_and_gradient(tn.code, tensors)
     # use Array to convert CuArray to CPU arrays
     return Array(logp)[], map(k->haskey(tn.fixedvertices, vars[k]) ? tn.fixedvertices[vars[k]] : argmax(grads[k]) - 1, 1:length(vars))
@@ -57,8 +57,8 @@ $(TYPEDSIGNATURES)
 
 Returns an output array containing largest log-probabilities.
 """
-function maximum_logp(tn::TensorNetworkModeling; usecuda=false)::AbstractArray{<:Tropical}
+function maximum_logp(tn::TensorNetworkModel; usecuda=false)::AbstractArray{<:Tropical}
     # generate tropical tensors with its elements being log(p).
-    tensors = map(t->Tropical.(log.(t)), generate_tensors(tn; usecuda, rescale=false))
+    tensors = map(t->Tropical.(log.(t)), adapt_tensors(tn; usecuda, rescale=false))
     return tn.code(tensors...)
 end
\ No newline at end of file
diff --git a/src/mmap.jl b/src/mmap.jl
index 5ccb5e0..3b8e614 100644
--- a/src/mmap.jl
+++ b/src/mmap.jl
@@ -20,10 +20,10 @@ Computing the most likely assignment to the query variables,  Xₘ ⊆ X after m
 * `vars` is the remaining (or not marginalized) degree of freedoms 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 [`TensorNetworkModeling`](@ref) instance for marginalizing certain variables.
+* `clusters` is the clusters, each element of this cluster is a [`TensorNetworkModel`](@ref) instance for marginalizing certain variables.
 * `fixedvertices` is a dictionary to specifiy degree of freedoms fixed to certain values, which should not have overlap with the marginalized variables.
 """
-struct MMAPModeling{LT,AT<:AbstractArray}
+struct MMAPModel{LT,AT<:AbstractArray}
     vars::Vector{LT}
     code::AbstractEinsum
     tensors::Vector{AT}
@@ -31,7 +31,7 @@ struct MMAPModeling{LT,AT<:AbstractArray}
     fixedvertices::Dict{LT,Int}
 end
 
-function Base.show(io::IO, mmap::MMAPModeling)
+function Base.show(io::IO, mmap::MMAPModel)
     open = getiyv(mmap.code)
     variables = join([string_var(var, open, mmap.fixedvertices) for var in mmap.vars], ", ")
     tc, sc, rw = timespacereadwrite_complexity(mmap)
@@ -40,17 +40,17 @@ function Base.show(io::IO, mmap::MMAPModeling)
     println(io, "marginalized variables: $(map(x->x.eliminated_vars, mmap.clusters))")
     print_tcscrw(io, tc, sc, rw)
 end
-Base.show(io::IO, ::MIME"text/plain", mmap::MMAPModeling) = Base.show(io, mmap)
+Base.show(io::IO, ::MIME"text/plain", mmap::MMAPModel) = Base.show(io, mmap)
 
 """
 $(TYPEDSIGNATURES)
 """
-get_vars(mmap::MMAPModeling) = mmap.vars
+get_vars(mmap::MMAPModel) = mmap.vars
 
 """
 $(TYPEDSIGNATURES)
 """
-function get_cards(mmap::MMAPModeling; fixedisone=false)
+function get_cards(mmap::MMAPModel; fixedisone=false)
     vars = get_vars(mmap)
     [fixedisone && haskey(mmap.fixedvertices, vars[k]) ? 1 : length(mmap.tensors[k]) for k=1:length(vars)]
 end
@@ -58,18 +58,18 @@ end
 """
 $(TYPEDSIGNATURES)
 """
-function MMAPModeling(instance::UAIInstance; marginalizedvertices, openvertices=(), optimizer=GreedyMethod(), simplifier=nothing)::MMAPModeling
-    return MMAPModeling(1:instance.nvars, instance.factors; marginalizedvertices, fixedvertices=Dict(zip(instance.obsvars, instance.obsvals .- 1)), optimizer, simplifier, openvertices)
+function MMAPModel(instance::UAIInstance; marginalizedvertices, openvertices=(), optimizer=GreedyMethod(), simplifier=nothing)::MMAPModel
+    return MMAPModel(1:instance.nvars, instance.factors; marginalizedvertices, fixedvertices=Dict(zip(instance.obsvars, instance.obsvals .- 1)), optimizer, simplifier, openvertices)
 end
 
 """
 $(TYPEDSIGNATURES)
 """
-function MMAPModeling(vars::AbstractVector{LT}, factors::Vector{<:Factor{T}}; marginalizedvertices, openvertices=(),
+function MMAPModel(vars::AbstractVector{LT}, factors::Vector{<:Factor{T}}; marginalizedvertices, openvertices=(),
             fixedvertices=Dict{LT,Int}(),
             optimizer=GreedyMethod(), simplifier=nothing,
             marginalize_optimizer=GreedyMethod(), marginalize_simplifier=nothing
-        )::MMAPModeling where {T,LT}
+        )::MMAPModel where {T,LT}
     all_ixs = [[[var] for var in vars]..., [[factor.vars...] for factor in factors]...]  # labels for vertex tensors (unity tensors) and edge tensors
     iy = collect(LT, openvertices)
     if !isempty(setdiff(iy, vars))
@@ -94,10 +94,10 @@ function MMAPModeling(vars::AbstractVector{LT}, factors::Vector{<:Factor{T}}; ma
     rem_indices = setdiff(1:length(all_ixs), vcat([c.second for c in subsets]...))
     remaining_tensors = all_tensors[rem_indices]
     code = optimize_code(EinCode([all_ixs[rem_indices]..., ixs...], iy), size_dict, optimizer, simplifier)
-    return MMAPModeling(setdiff(vars, marginalizedvertices), code, remaining_tensors, clusters, fixedvertices)
+    return MMAPModel(setdiff(vars, marginalizedvertices), code, remaining_tensors, clusters, fixedvertices)
 end
 
-function OMEinsum.timespacereadwrite_complexity(mmap::MMAPModeling{LT}) where LT
+function OMEinsum.timespacereadwrite_complexity(mmap::MMAPModel{LT}) where LT
     # extract size
     size_dict = Dict(zip(get_vars(mmap), get_cards(mmap; fixedisone=true)))
     sc = -Inf
@@ -120,10 +120,10 @@ function OMEinsum.timespacereadwrite_complexity(mmap::MMAPModeling{LT}) where LT
     push!(rws, tc)
     OMEinsum.OMEinsumContractionOrders.log2sumexp2(tcs), max(sc, sci), OMEinsum.OMEinsumContractionOrders.log2sumexp2(rws)
 end
-OMEinsum.timespace_complexity(mmap::MMAPModeling) = timespacereadwrite_complexity(mmap)[1:2]
+OMEinsum.timespace_complexity(mmap::MMAPModel) = timespacereadwrite_complexity(mmap)[1:2]
 
-function generate_tensors(mmap::MMAPModeling; usecuda, rescale)
-    return [generate_tensors(mmap.code, mmap.tensors, mmap.fixedvertices; usecuda, rescale)...,
+function adapt_tensors(mmap::MMAPModel; usecuda, rescale)
+    return [adapt_tensors(mmap.code, mmap.tensors, mmap.fixedvertices; usecuda, rescale)...,
         map(cluster->probability(cluster; fixedvertices=mmap.fixedvertices, usecuda, rescale), mmap.clusters)...]
 end
 
@@ -175,9 +175,9 @@ end
 """
 $(TYPEDSIGNATURES)
 """
-function most_probable_config(mmap::MMAPModeling; usecuda=false)::Tuple{Tropical,Vector}
+function most_probable_config(mmap::MMAPModel; usecuda=false)::Tuple{Tropical,Vector}
     vars = get_vars(mmap)
-    tensors = map(t->OMEinsum.asarray(Tropical.(log.(t)), t), generate_tensors(mmap; usecuda, rescale=false))
+    tensors = map(t->OMEinsum.asarray(Tropical.(log.(t)), t), adapt_tensors(mmap; usecuda, rescale=false))
     logp, grads = cost_and_gradient(mmap.code, tensors)
     # use Array to convert CuArray to CPU arrays
     return Array(logp)[], map(k->haskey(mmap.fixedvertices, vars[k]) ? mmap.fixedvertices[vars[k]] : argmax(grads[k]) - 1, 1:length(vars))
@@ -186,26 +186,26 @@ end
 """
 $(TYPEDSIGNATURES)
 """
-function maximum_logp(mmap::MMAPModeling; usecuda=false)::AbstractArray{<:Tropical}
-    tensors = map(t->OMEinsum.asarray(Tropical.(log.(t)), t), generate_tensors(mmap; usecuda, rescale=false))
+function maximum_logp(mmap::MMAPModel; usecuda=false)::AbstractArray{<:Tropical}
+    tensors = map(t->OMEinsum.asarray(Tropical.(log.(t)), t), adapt_tensors(mmap; usecuda, rescale=false))
     return mmap.code(tensors...)
 end
 
 """
 $(TYPEDSIGNATURES)
 """
-function log_probability(mmap::MMAPModeling, config::Union{Dict, AbstractVector}; rescale=true, usecuda=false)::Real
+function log_probability(mmap::MMAPModel, config::Union{Dict, AbstractVector}; rescale=true, usecuda=false)::Real
     @assert length(get_vars(mmap)) == length(config)
     fixedvertices =  config isa AbstractVector ? Dict(zip(get_vars(mmap), config)) : config
     assign = merge(mmap.fixedvertices, fixedvertices)
     # two contributions to the probability, not-clustered tensors and clusters.
     m1 = sum(x->log(x[2][(getindex.(Ref(assign), x[1]) .+ 1)...]), zip(getixsv(mmap.code), mmap.tensors))
-    m2 = sum(cluster->probability(cluster; fixedvertices, usecuda, rescale).logf, mmap.clusters)
+    m2 = sum(cluster->probability(cluster; fixedvertices, usecuda, rescale).log_factor, mmap.clusters)
     return m1 + m2
 end
 
 function probability(c::Cluster; fixedvertices, usecuda, rescale)::AbstractArray
-    tensors = generate_tensors(c.code, c.tensors, fixedvertices; usecuda, rescale)
+    tensors = adapt_tensors(c.code, c.tensors, fixedvertices; usecuda, rescale)
     return c.code(tensors...)
 end
 
diff --git a/test/cuda.jl b/test/cuda.jl
index 8c0e38e..5f2ec96 100644
--- a/test/cuda.jl
+++ b/test/cuda.jl
@@ -8,7 +8,7 @@ CUDA.allowscalar(false)
     instance = read_uai_problem("Promedus_14")
 
     # does not optimize over open vertices
-    tn = TensorNetworkModeling(instance; optimizer=TreeSA(ntrials=1, niters=2, βs=1:0.1:40))
+    tn = TensorNetworkModel(instance; optimizer=TreeSA(ntrials=1, niters=2, βs=1:0.1:40))
     @info contraction_complexity(tn)
     @time marginals2 = marginals(tn; usecuda=true)
     @test all(x->x isa CuArray, marginals2)
@@ -25,7 +25,7 @@ end
     instance = read_uai_problem("Promedus_14")
 
     # does not optimize over open vertices
-    tn = TensorNetworkModeling(instance; optimizer=TreeSA(ntrials=1, niters=2, βs=1:0.1:40))
+    tn = TensorNetworkModel(instance; optimizer=TreeSA(ntrials=1, niters=2, βs=1:0.1:40))
     @info contraction_complexity(tn)
     most_probable_config(tn)
     @time logp, config = most_probable_config(tn; usecuda=true)
@@ -40,16 +40,16 @@ end
     instance = read_uai_problem("Promedus_14")
 
     optimizer=TreeSA(ntrials=1, niters=2, βs=1:0.1:40)
-    tn_ref = TensorNetworkModeling(instance; optimizer)
+    tn_ref = TensorNetworkModel(instance; optimizer)
     # does not marginalize any var
-    tn = MMAPModeling(instance; marginalizedvertices=Int[], optimizer)
+    tn = MMAPModel(instance; marginalizedvertices=Int[], optimizer)
     r1, r2 = maximum_logp(tn_ref; usecuda=true), maximum_logp(tn; usecuda=true)
     @test r1 isa CuArray
     @test r2 isa CuArray
     @test r1 ≈ r2
 
     # marginalize all vars
-    tn2 = MMAPModeling(instance; marginalizedvertices=collect(1:instance.nvars), optimizer)
+    tn2 = MMAPModel(instance; marginalizedvertices=collect(1:instance.nvars), optimizer)
     cup = probability(tn_ref; usecuda=true)
     culogp = maximum_logp(tn2; usecuda=true)
     @test cup isa RescaledArray{T, N, <:CuArray} where {T, N}
@@ -57,7 +57,7 @@ end
     @test Array(cup)[] ≈ exp(Array(culogp)[].n)
 
     # does not optimize over open vertices
-    tn3 = MMAPModeling(instance; marginalizedvertices=[2,4,6], optimizer)
+    tn3 = MMAPModel(instance; marginalizedvertices=[2,4,6], optimizer)
     logp, config = most_probable_config(tn3; usecuda=true)
     @test log_probability(tn3, config) ≈ logp.n
 end 
diff --git a/test/inference.jl b/test/inference.jl
index cccb0ff..470a844 100644
--- a/test/inference.jl
+++ b/test/inference.jl
@@ -15,13 +15,13 @@ end
 @testset "cached, rescaled contract" begin
     problem = read_uai_problem("Promedus_14")
     optimizer = TreeSA(ntrials=1, niters=5, βs=0.1:0.1:100)
-    tn = TensorNetworkModeling(problem; optimizer)
+    tn = TensorNetworkModel(problem; optimizer)
     p1 = probability(tn; usecuda=false, rescale=false)
     p2 = probability(tn; usecuda=false, rescale=true)
     @test p1 ≈ Array(p2)
     
     # cached contract
-    xs = TensorInference.generate_tensors(tn; usecuda=false, rescale=true)
+    xs = TensorInference.adapt_tensors(tn; usecuda=false, rescale=true)
     size_dict = OMEinsum.get_size_dict!(getixsv(tn.code), xs, Dict{Int,Int}())
     cache = TensorInference.cached_einsum(tn.code, xs, size_dict)
     @test cache.content isa RescaledArray
@@ -73,7 +73,7 @@ end
             problem = read_uai_problem(problem)
 
             # does not optimize over open vertices
-            tn = TensorNetworkModeling(problem; optimizer)
+            tn = TensorNetworkModel(problem; optimizer)
             sc = contraction_complexity(tn).sc
             if sc > 28
               error("space complexity too large! got $(sc)")
diff --git a/test/maxprob.jl b/test/maxprob.jl
index 0f5fafd..81d495d 100644
--- a/test/maxprob.jl
+++ b/test/maxprob.jl
@@ -7,7 +7,7 @@ using TensorInference
     instance = read_uai_problem("Promedus_14")
 
     # does not optimize over open vertices
-    tn = TensorNetworkModeling(instance; optimizer=TreeSA(ntrials=1, niters=2, βs=1:0.1:40))
+    tn = TensorNetworkModel(instance; optimizer=TreeSA(ntrials=1, niters=2, βs=1:0.1:40))
     @info contraction_complexity(tn)
     most_probable_config(tn)
     @time logp, config = most_probable_config(tn)
diff --git a/test/mmap.jl b/test/mmap.jl
index 5a38590..31dbdc0 100644
--- a/test/mmap.jl
+++ b/test/mmap.jl
@@ -12,19 +12,19 @@ end
     instance = read_uai_problem("Promedus_14")
 
     optimizer=TreeSA(ntrials=1, niters=2, βs=1:0.1:40)
-    tn_ref = TensorNetworkModeling(instance; optimizer)
+    tn_ref = TensorNetworkModel(instance; optimizer)
     # does not marginalize any var
-    mmap = MMAPModeling(instance; marginalizedvertices=Int[], optimizer)
+    mmap = MMAPModel(instance; marginalizedvertices=Int[], optimizer)
     @info(mmap)
     @test maximum_logp(tn_ref) ≈ maximum_logp(mmap)
 
     # marginalize all vars
-    mmap2 = MMAPModeling(instance; marginalizedvertices=collect(1:instance.nvars), optimizer)
+    mmap2 = MMAPModel(instance; marginalizedvertices=collect(1:instance.nvars), optimizer)
     @info(mmap2)
     @test Array(probability(tn_ref))[] ≈ exp(maximum_logp(mmap2)[].n)
 
     # does not optimize over open vertices
-    mmap3 = MMAPModeling(instance; marginalizedvertices=[2,4,6], optimizer)
+    mmap3 = MMAPModel(instance; marginalizedvertices=[2,4,6], optimizer)
     @info(mmap3)
     logp, config = most_probable_config(mmap3)
     @test log_probability(mmap3, config) ≈ logp.n

From c4391ec9e3af4e24475a4d7b83e5a7ed784aa22b Mon Sep 17 00:00:00 2001
From: GiggleLiu <cacate0129@gmail.com>
Date: Thu, 3 Nov 2022 22:24:49 +0800
Subject: [PATCH 6/6] fix a bug

---
 src/inference.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/inference.jl b/src/inference.jl
index 804edbf..acf9836 100644
--- a/src/inference.jl
+++ b/src/inference.jl
@@ -132,7 +132,7 @@ function marginals(tn::TensorNetworkModel; usecuda=false, rescale=true)::Vector
     cost, grads = cost_and_gradient(tn.code, adapt_tensors(tn; usecuda, rescale))
     @debug "cost = $cost"
     if rescale
-        return LinearAlgebra.normalize!.(getfield.(grads[1:length(vars)], :rescaled_value), 1)
+        return LinearAlgebra.normalize!.(getfield.(grads[1:length(vars)], :normalized_value), 1)
     else
         return LinearAlgebra.normalize!.(grads[1:length(vars)], 1)
     end