fix cuda tests

Author: GiggleLiu

NEWS.md

@@ -83,7 +83,6 @@ project_relative_path(xs...) = normpath(joinpath(dirname(dirname(pathof(@__MODUL
 include("Mods.jl/src/Mods.jl")
 using .Mods
 
-include("utils.jl")
 include("arithematics.jl")
 include("networks.jl")
 include("graph_polynomials.jl")

docs/src/tensornetwork.md

@@ -158,7 +158,7 @@ function solution_ad(code::Union{NestedEinsum,SlicedEinsum}, @nospecialize(xsa),
     @debug "generating masked tree..."
     mt = generate_masktree(SingleConfig(), code, c, ymask, size_dict)
     config = read_config!(code, mt, Dict())
-    if length(config) !== length(labels(code))  # equal to the # of degree of freedoms
+    if length(config) !== length(uniquelabels(code))  # equal to the # of degree of freedoms
         error("configuration `$(config)` is not fully determined!")
     end
     n, config

src/GenericTensorNetworks.jl

@@ -11,85 +11,70 @@ function config_type(::Type{T}, n, num_flavors; all::Bool, tree_storage::Bool) w
 end
 
 """
-    largest_solutions(problem; all=false, usecuda=false, invert=false, tree_storage::Bool=false)
+    largest_solutions(net::GenericTensorNetwork; all=false, usecuda=false, invert=false, tree_storage::Bool=false, T=Float64)
     
-Find optimal solutions with bounding.
-
-* When `all` is true, the program will use set for enumerate all possible solutions, otherwise, it will return one solution for each size.
-* `usecuda` can not be true if you want to use set to enumerate all possible solutions.
-* If `invert` is true, find the minimum.
-* If `tree_storage` is true, use [`SumProductTree`](@ref) as the storage of solutions.
+Find optimal solutions, with bounding. Please check [`solutions`](@ref) for argument descriptions.
 """
-function largest_solutions(gp::GenericTensorNetwork; all=false, usecuda=false, invert=false, tree_storage::Bool=false, T=Float64)
+function largest_solutions(net::GenericTensorNetwork; all=false, usecuda=false, invert=false, tree_storage::Bool=false, T=Float64)
     if all && usecuda
         throw(ArgumentError("ConfigEnumerator can not be computed on GPU!"))
     end
-    xst = generate_tensors(_x(Tropical{T}; invert), gp)
-    ymask = trues(fill(num_flavors(gp), length(getiyv(gp.code)))...)
+    xst = generate_tensors(_x(Tropical{T}; invert), net)
+    ymask = trues(fill(num_flavors(net), length(getiyv(net.code)))...)
     if usecuda
         xst = togpu.(xst)
         ymask = togpu(ymask)
     end
     if all
         # we use `Float64` as default because we want to support weighted graphs.
-        T = config_type(CountingTropical{T,T}, length(labels(gp)), num_flavors(gp); all, tree_storage)
-        xs = generate_tensors(_x(T; invert), gp)
-        ret = bounding_contract(AllConfigs{1}(), gp.code, xst, ymask, xs)
+        T = config_type(CountingTropical{T,T}, length(variables(net)), num_flavors(net); all, tree_storage)
+        xs = generate_tensors(_x(T; invert), net)
+        ret = bounding_contract(AllConfigs{1}(), net.code, xst, ymask, xs)
         return invert ? asarray(post_invert_exponent.(ret), ret) : ret
     else
         @assert ndims(ymask) == 0
-        t, res = solution_ad(gp.code, xst, ymask)
+        t, res = solution_ad(net.code, xst, ymask)
         ret = fill(CountingTropical(asscalar(t).n, ConfigSampler(StaticBitVector(map(l->res[l], 1:length(res))))))
         return invert ? asarray(post_invert_exponent.(ret), ret) : ret
     end
 end
 
 """
-    solutions(problem, basetype; all, usecuda=false, invert=false, tree_storage::Bool=false)
+    solutions(net::GenericTensorNetwork, ::Type{BT}; all::Bool, usecuda::Bool=false, invert::Bool=false, tree_storage::Bool=false) where BT
     
-General routine to find solutions without bounding,
+Find all solutions, solutions with largest sizes or solutions with smallest sizes. Bounding is not supported.
 
-* `basetype` can be a type with counting field,
+### Arguments
+- `net` is a [`GenericTensorNetwork`](@ref) instance.
+- `BT` is the data types used for computing, which can be
     * `CountingTropical{Float64,Float64}` for finding optimal solutions,
     * `Polynomial{Float64, :x}` for enumerating all solutions,
     * `Max2Poly{Float64,Float64}` for optimal and suboptimal solutions.
-* When `all` is true, the program will use set for enumerate all possible solutions, otherwise, it will return one solution for each size.
-* `usecuda` can not be true if you want to use set to enumerate all possible solutions.
-* If `tree_storage` is true, use [`SumProductTree`](@ref) as the storage of solutions.
+
+### Keyword arguments
+- `all` is an indicator whether to find all solutions or just one of them.
+- `usecuda` is an indicator of using CUDA or not, which must be false if `all` is true.
+- `invert` is an indicator of whether flip the size or not. If true, instead of finding the maximum, it find the minimum.
+- `tree_storage` is an indicator of whether using more compact [`SumProductTree`](@ref) as the storage or not.
 """
-function solutions(gp::GenericTensorNetwork, ::Type{BT}; all::Bool, usecuda::Bool=false, invert::Bool=false, tree_storage::Bool=false) where BT
+function solutions(net::GenericTensorNetwork, ::Type{BT}; all::Bool, usecuda::Bool=false, invert::Bool=false, tree_storage::Bool=false) where BT
     if all && usecuda
         throw(ArgumentError("ConfigEnumerator can not be computed on GPU!"))
     end
-    T = config_type(BT, length(labels(gp)), num_flavors(gp); all, tree_storage)
-    ret = contractx(gp, _x(T; invert); usecuda=usecuda)
+    T = config_type(BT, length(variables(net)), num_flavors(net); all, tree_storage)
+    ret = contractx(net, _x(T; invert); usecuda=usecuda)
     return invert ? asarray(post_invert_exponent.(ret), ret) : ret
 end
 
-"""
-    largest2_solutions(problem; all=true, usecuda=false, invert=false, tree_storage::Bool=false)
-
-Finding optimal and suboptimal solutions.
-"""
-largest2_solutions(gp::GenericTensorNetwork; all=true, usecuda=false, invert::Bool=false, T=Float64) = solutions(gp, Max2Poly{T,T}; all, usecuda, invert)
-
-function largestk_solutions(gp::GenericTensorNetwork, k::Int; invert::Bool=false, tree_storage::Bool=false, T=Float64)
-    xst = generate_tensors(_x(Tropical{T}; invert), gp)
-    ymask = trues(fill(2, length(getiyv(gp.code)))...)
-    T = config_type(TruncatedPoly{k,T,T}, length(labels(gp)), num_flavors(gp); all=true, tree_storage)
-    xs = generate_tensors(_x(T; invert), gp)
-    ret = bounding_contract(AllConfigs{k}(), gp.code, xst, ymask, xs)
+function largestk_solutions(net::GenericTensorNetwork, k::Int; invert::Bool=false, tree_storage::Bool=false, T=Float64)
+    xst = generate_tensors(_x(Tropical{T}; invert), net)
+    ymask = trues(fill(2, length(getiyv(net.code)))...)
+    T = config_type(TruncatedPoly{k,T,T}, length(variables(net)), num_flavors(net); all=true, tree_storage)
+    xs = generate_tensors(_x(T; invert), net)
+    ret = bounding_contract(AllConfigs{k}(), net.code, xst, ymask, xs)
     return invert ? asarray(post_invert_exponent.(ret), ret) : ret
 end
 
-"""
-    all_solutions(problem)
-
-Finding all solutions grouped by size.
-e.g. when the problem is [`MaximalIS`](@ref), it computes all maximal independent sets, or the maximal cliques of it complement.
-"""
-all_solutions(gp::GenericTensorNetwork; T=Float64) = solutions(gp, Polynomial{T,:x}, all=true, usecuda=false, tree_storage=false)
-
 # NOTE: do we have more efficient way to compute it?
 # NOTE: doing pair-wise Hamming distance might be biased?
 """

src/bounding.jl

@@ -9,4 +9,13 @@
 @deprecate get_weights(args...; kwargs...) weights(args...; kwargs...)
 @deprecate chweights(args...; kwargs...) set_weights(args...; kwargs...)
 
-@deprecate GraphProblem() ConstraintSatisfactionProblem()
+@deprecate spinglass_energy(args...; kwargs...) energy(args...; kwargs...)
+@deprecate unit_disk_graph(args...; kwargs...) UnitDiskGraph(args...; kwargs...)
+@deprecate solve(problem::ConstraintSatisfactionProblem, args...; kwargs...) solve(GenericTensorNetwork(problem), args...; kwargs...)
+
+"""
+const GraphProblem = ConstraintSatisfactionProblem
+
+Deprecated. Use `ConstraintSatisfactionProblem` instead.
+"""
+const GraphProblem = ConstraintSatisfactionProblem

src/configurations.jl

@@ -390,37 +390,37 @@ Memory estimation in number of bytes to compute certain `property` of a `problem
 `T` is the base type.
 """
 function estimate_memory(problem::GenericTensorNetwork, property::AbstractProperty; T=Float64)::Real
-    _estimate_memory(tensor_element_type(T, length(labels(problem)), num_flavors(problem), property), problem)
+    _estimate_memory(tensor_element_type(T, length(variables(problem)), num_flavors(problem), property), problem)
 end
 function estimate_memory(problem::GenericTensorNetwork, property::Union{SingleConfigMax{K,BOUNDED},SingleConfigMin{K,BOUNDED}}; T=Float64) where {K, BOUNDED}
     tc, sc, rw = contraction_complexity(problem.code, _size_dict(problem))
     # caching all tensors is equivalent to counting the total number of writes
     if K === Single && BOUNDED
         return ceil(Int, exp2(rw - 1)) * sizeof(Tropical{T})
     elseif K === Single && !BOUNDED
-        n, nf = length(labels(problem)), num_flavors(problem)
+        n, nf = length(variables(problem)), num_flavors(problem)
         return peak_memory(problem.code, _size_dict(problem)) * (sizeof(tensor_element_type(T, n, nf, property)))
     else
         # NOTE: the integer `K` case does not respect bounding
-        n, nf = length(labels(problem)), num_flavors(problem)
+        n, nf = length(variables(problem)), num_flavors(problem)
         TT = tensor_element_type(T, n, nf, property)
         return peak_memory(problem.code, _size_dict(problem)) * (sizeof(tensor_element_type(T, n, nf, SingleConfigMax{Single,BOUNDED}())) * K + sizeof(TT))
     end
 end
 function estimate_memory(problem::GenericTensorNetwork, ::GraphPolynomial{:polynomial}; T=Float64)
     # this is the upper bound
-    return peak_memory(problem.code, _size_dict(problem)) * (sizeof(T) * length(labels(problem)))
+    return peak_memory(problem.code, _size_dict(problem)) * (sizeof(T) * length(variables(problem)))
 end
 function estimate_memory(problem::GenericTensorNetwork, ::GraphPolynomial{:laurent}; T=Float64)
     # this is the upper bound
-    return peak_memory(problem.code, _size_dict(problem)) * (sizeof(T) * length(labels(problem)))
+    return peak_memory(problem.code, _size_dict(problem)) * (sizeof(T) * length(variables(problem)))
 end
 function estimate_memory(problem::GenericTensorNetwork, ::Union{SizeMax{K},SizeMin{K}}; T=Float64) where K
     return peak_memory(problem.code, _size_dict(problem)) * (sizeof(T) * _asint(K))
 end
 
 function _size_dict(problem)
-    lbs = labels(problem)
+    lbs = variables(problem)
     nf = num_flavors(problem)
     return Dict([lb=>nf for lb in lbs])
 end

src/deprecate.jl

@@ -32,6 +32,10 @@ function GenericTensorNetwork(problem::ConstraintSatisfactionProblem; openvertic
     code = _optimize_code(rcode, uniformsize_fix(rcode, num_flavors(problem), fixedvertices), optimizer, MergeVectors())
     return GenericTensorNetwork(problem, code, Dict{labeltype(code),Int}(fixedvertices))
 end
+# a unified interface to optimize the contraction code
+_optimize_code(code, size_dict, optimizer::Nothing, simplifier) = code
+_optimize_code(code, size_dict, optimizer, simplifier) = optimize_code(code, size_dict, optimizer, simplifier)
+
 function Base.show(io::IO, tn::GenericTensorNetwork)
     println(io, "$(typeof(tn))")
     println(io, "- open vertices: $(getiyv(tn.code))")
@@ -184,6 +188,9 @@ function _pow(x::LaurentPolynomial{BS,X}, i) where {BS,X}
 end
 
 ####### Extra utilities #######
+"""Upload a tensor network to GPU, need `using CUDA` to activate this extension."""
+function togpu end
+
 """
     mis_compactify!(tropicaltensor; potential=nothing)
 

src/interfaces.jl

@@ -1,7 +1,7 @@
 using GenericTensorNetworks, Test, Graphs
 using OMEinsum
 using TropicalNumbers: CountingTropicalF64
-using GenericTensorNetworks: _onehotv, _x, sampler_type, set_type, largest_solutions, largest2_solutions, solutions, all_solutions, largestk_solutions, AllConfigs, SingleConfig, max_size, max_size_count
+using GenericTensorNetworks: _onehotv, _x, sampler_type, set_type, largest_solutions, solutions, largestk_solutions, AllConfigs, SingleConfig, max_size, max_size_count
 using GenericTensorNetworks: graph_polynomial
 
 @testset "Config types" begin
@@ -48,14 +48,11 @@ end
         res5 = largest_solutions(code; all=false)[]
         @test res5.n == res0
         @test res5.c.data ∈ res2.c.data
-        res6 = largest2_solutions(code; all=true)[]
-        res6_ = largestk_solutions(code, 2)[]
-        res7 = all_solutions(code)[]
+        res6 = largestk_solutions(code, 2)[]
+        res7 = GenericTensorNetworks.solutions(code, Polynomial{Float64,:x}, all=true, usecuda=false, tree_storage=false)[]
         idp = graph_polynomial(code, Val(:finitefield))[]
         @test all(x->x ∈ res7.coeffs[end-1].data, res6.coeffs[1].data)
         @test all(x->x ∈ res7.coeffs[end].data, res6.coeffs[2].data)
-        @test all(x->x ∈ res7.coeffs[end-1].data, res6_.coeffs[1].data)
-        @test all(x->x ∈ res7.coeffs[end].data, res6_.coeffs[2].data)
         for (i, (s, c)) in enumerate(zip(res7.coeffs, idp.coeffs))
             @test length(s) == c
             @test all(x->count_ones(x)==(i-1), s.data)