Switch from Requires to extensions

Project.toml

@@ -5,24 +5,28 @@ version = "0.4.0"
 
 [deps]
 Artifacts = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"
-CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
-GenericTensorNetworks = "3521c873-ad32-4bb4-b63d-f4f178f42b49"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 OMEinsum = "ebe7aa44-baf0-506c-a96f-8464559b3922"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
-Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 TropicalNumbers = "b3a74e9c-7526-4576-a4eb-79c0d4c32334"
 
+[weakdeps]
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+GenericTensorNetworks = "3521c873-ad32-4bb4-b63d-f4f178f42b49"
+
+[extensions]
+TensorInferenceCUDAExt = "CUDA"
+TensorInferenceGTNExt = "GenericTensorNetworks"
+
 [compat]
 CUDA = "4, 5"
 DocStringExtensions = "0.8.6, 0.9"
 GenericTensorNetworks = "1"
 OMEinsum = "0.7"
 PrecompileTools = "1"
-Requires = "1"
 StatsBase = "0.34"
 TropicalNumbers = "0.5.4, 0.6"
-julia = "1.3"
+julia = "1.9"

src/cuda.jl → ext/TensorInferenceCUDAExt.jl

@@ -1,4 +1,7 @@
-using .CUDA: CuArray
+module TensorInferenceCUDAExt
+using CUDA: CuArray
+import CUDA
+import TensorInference: match_arraytype, keep_only!, onehot_like, togpu
 
 function onehot_like(A::CuArray, j)
     mask = zero(A)
@@ -15,3 +18,7 @@
     CUDA.@allowscalar x[j] = hotvalue
     return x
 end
+
+togpu(x::AbstractArray) = CuArray(x)
+
+end

src/generictensornetworks.jl → ext/TensorInferenceGTNExt.jl

@@ -1,7 +1,7 @@
-using .GenericTensorNetworks: generate_tensors, GraphProblem, flavors, labels
-
-# update models
-export update_temperature
+module TensorInferenceGTNExt
+using TensorInference, TensorInference.OMEinsum
+using TensorInference: TYPEDSIGNATURES, Factor
+using GenericTensorNetworks: generate_tensors, GraphProblem, flavors, labels
 
 """
 $TYPEDSIGNATURES
@@ -24,18 +24,7 @@ function TensorInference.TensorNetworkModel(problem::GraphProblem, β::Real; evi
 	return TensorNetworkModel(lbs, fill(nflavors, length(lbs)), factors; openvars=iy, evidence, optimizer, simplifier, mars)
 end
 
-"""
-$TYPEDSIGNATURES
-
-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`](@ref) instance.
-- `problem` is the target constraint satisfiability problem.
-- `β` is the inverse temperature.
-"""
-function update_temperature(tnet::TensorNetworkModel, problem::GraphProblem, β::Real)
+function TensorInference.update_temperature(tnet::TensorNetworkModel, problem::GraphProblem, β::Real)
 	tensors = generate_tensors(exp(β), problem)
     alltensors = [tnet.tensors[1:end-length(tensors)]..., tensors...]
     return TensorNetworkModel(tnet.vars, tnet.code, alltensors, tnet.evidence, tnet.mars)
@@ -64,4 +53,5 @@ It is about one or two hours of works. If you need it, please file an issue to l
 end
 
 @info "`TensorInference` loaded `GenericTensorNetworks` extension successfully,
-`TensorNetworkModel` and `MMAPModel` can be used for converting a `GraphProblem` to a probabilistic model now."
+`TensorNetworkModel` and `MMAPModel` can be used for converting a `GraphProblem` to a probabilistic model now."
+end

src/TensorInference.jl

@@ -34,6 +34,9 @@ export sample
 # MMAP
 export MMAPModel
 
+# for GenericTensorNetworks
+export update_temperature
+
 include("Core.jl")
 include("RescaledArray.jl")
 include("utils.jl")
@@ -42,12 +45,6 @@ include("map.jl")
 include("mmap.jl")
 include("sampling.jl")
 
-using Requires
-function __init__()
-    @require CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" include("cuda.jl")
-    @require GenericTensorNetworks = "3521c873-ad32-4bb4-b63d-f4f178f42b49" include("generictensornetworks.jl")
-end
-
 # import PrecompileTools
 # PrecompileTools.@setup_workload begin
 #     # Putting some things in `@setup_workload` instead of `@compile_workload` can reduce the size of the
@@ -57,4 +54,17 @@ end
 #     end
 # end
 
+"""
+$TYPEDSIGNATURES
+
+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`](@ref) instance.
+- `problem` is the target constraint satisfiability problem.
+- `β` is the inverse temperature.
+"""
+function update_temperature end
+
 end # module

src/mar.jl

@@ -7,7 +7,7 @@ function adapt_tensors(code, tensors, evidence; usecuda, rescale)
     map(tensors, ixs) do t, ix
         dims = map(ixi -> ixi ∉ keys(evidence) ? Colon() : ((evidence[ixi] + 1):(evidence[ixi] + 1)), ix)
         t2 = t[dims...]
-        t3 = usecuda ? CuArray(t2) : t2
+        t3 = usecuda ? togpu(t2) : t2
         rescale ? rescale_array(t3) : t3
     end
 end

src/utils.jl

@@ -305,3 +305,5 @@
     artifact_hash = Pkg.Artifacts.artifact_hash(artifact_name, artifact_toml)
     return Pkg.Artifacts.artifact_path(artifact_hash)
 end
+
+togpu(x) = error("You must import CUDA with `using CUDA` before using GPU!")