Create CUDA and GenericTensorNetworks extensions

Author: GiggleLiu

Project.toml

@@ -5,9 +5,7 @@ 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"
@@ -16,13 +14,21 @@ 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"
+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"
-julia = "1.3"
+TropicalNumbers = "0.5.4, 0.6"
+julia = "1.9"

src/cuda.jl ext/TensorInferenceCUDAExt.jlsrc/cuda.jl renamed to 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 @@ function keep_only!(x::CuArray{T}, j) where T
     CUDA.@allowscalar x[j] = hotvalue
     return x
 end
+
+togpu(x::AbstractArray) = CuArray(x)
+
+end

src/generictensornetworks.jl ext/TensorInferenceGTNExt.jlsrc/generictensornetworks.jl renamed to ext/TensorInferenceGTNExt.jl

@@ -1,7 +1,8 @@
-using .GenericTensorNetworks: generate_tensors, GraphProblem, flavors, labels
-
-# update models
-export update_temperature
+module TensorInferenceGTNExt
+using TensorInference, TensorInference.OMEinsum
+using TensorInference: TYPEDSIGNATURES, Factor
+import TensorInference: update_temperature
+using GenericTensorNetworks: generate_tensors, GraphProblem, flavors, labels
 
 """
 $TYPEDSIGNATURES
@@ -64,4 +65,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,10 @@ export sample
 # MMAP
 export MMAPModel
 
+# for GenericTensorNetworks
+export update_temperature
+function update_temperature end
+
 include("Core.jl")
 include("RescaledArray.jl")
 include("utils.jl")
@@ -42,12 +46,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

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 @@ function get_artifact_path(artifact_name::String)
     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!")