1.9 extensions (#52)

* 1.9 extensions, first commit * load Requires before __init__ * typo in inports * Apply suggestions from code review * more qualifiers in JuMP ext * typos * use LinearAlgebra in optim ext, change some things in jump ext * compats * using Parameters in JuMP ext * compat for GPUArrays * add weakdeps to extras * use LinearAlgebra in JuMP ext * add Statistics to JuMP ext * typo in GPUArrays ext * typos found while testing for ChainRules and JLArrays --------- Co-authored-by: Júlio Hoffimann <[email protected]>
JuliaML · May 17, 2023 · 30d39c2 · 30d39c2
1 parent 7d2aa83
commit 30d39c2
Show file tree

Hide file tree

Showing 14 changed files with 183 additions and 50 deletions.
diff --git a/Project.toml b/Project.toml
@@ -11,8 +11,40 @@ Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 
+[weakdeps]
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+Convex = "f65535da-76fb-5f13-bab9-19810c17039a"
+ECOS = "e2685f51-7e38-5353-a97d-a921fd2c8199"
+GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
+Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9"
+JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
+Optim = "429524aa-4258-5aef-a3af-852621145aeb"
+
+[extensions]
+DensityRatioEstimationChainRulesCoreExt = "ChainRulesCore"
+DensityRatioEstimationGPUArraysExt = "GPUArrays"
+DensityRatioEstimationOptimExt = "Optim"
+DensityRatioEstimationConvexExt = ["Convex", "ECOS"]
+DensityRatioEstimationJuMPExt = ["JuMP", "Ipopt"]
+
 [compat]
+ChainRulesCore = "1"
+Convex = "0.15"
+ECOS = "1"
+GPUArrays = "8"
+Ipopt = "1"
+JuMP = "1"
+Optim = "1"
 Parameters = "0.12"
 Requires = "1"
 StatsBase = "0.32, 0.33"
 julia = "1.6"
+
+[extras]
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+Convex = "f65535da-76fb-5f13-bab9-19810c17039a"
+ECOS = "e2685f51-7e38-5353-a97d-a921fd2c8199"
+GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
+Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9"
+JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
+Optim = "429524aa-4258-5aef-a3af-852621145aeb"
diff --git a/ext/DensityRatioEstimationChainRulesCoreExt.jl b/ext/DensityRatioEstimationChainRulesCoreExt.jl
@@ -0,0 +1,16 @@
+# ------------------------------------------------------------------
+# Licensed under the MIT License. See LICENSE in the project root.
+# ------------------------------------------------------------------
+module DensityRatioEstimationChainRulesCoreExt
+
+if isdefined(Base, :get_extension)
+    using DensityRatioEstimation
+    using ChainRulesCore
+else
+    using ..DensityRatioEstimation
+    using ..ChainRulesCore
+end
+
+ChainRulesCore.@non_differentiable DensityRatioEstimation.safe_diagm(::Any,::Any)
+
+end #module
diff --git a/ext/DensityRatioEstimationConvexExt.jl b/ext/DensityRatioEstimationConvexExt.jl
@@ -0,0 +1,20 @@
+# ------------------------------------------------------------------
+# Licensed under the MIT License. See LICENSE in the project root.
+# ------------------------------------------------------------------
+module DensityRatioEstimationConvexExt
+
+    if isdefined(Base, :get_extension)
+        using DensityRatioEstimation
+        using DensityRatioEstimation: KLIEP, ConvexLib
+        using Convex
+        using ECOS
+    else
+        using ..DensityRatioEstimation
+        using ..DensityRatioEstimation: KLIEP, ConvexLib
+        using ..Convex
+        using ..ECOS
+    end
+
+    include("../src/kliep/convex.jl")
+
+end #module
diff --git a/ext/DensityRatioEstimationGPUArraysExt.jl b/ext/DensityRatioEstimationGPUArraysExt.jl
@@ -0,0 +1,22 @@
+# ------------------------------------------------------------------
+# Licensed under the MIT License. See LICENSE in the project root.
+# ------------------------------------------------------------------
+module DensityRatioEstimationGPUArraysExt
+
+    if isdefined(Base, :get_extension)
+        using DensityRatioEstimation
+        using GPUArrays
+    else
+        using ..DensityRatioEstimation
+        using ..GPUArrays
+    end
+    using LinearAlgebra
+
+    # Aviod `mat + a * I` with CUDA which involes scalar operations and is slow
+    function DensityRatioEstimation.safe_diagm(mat::M, a::T) where M<:GPUArrays.AbstractGPUArray{T, 2} where T
+        diag = similar(mat,size(m,1))
+        fill!(diag,a)
+        LinearAlgebra.Diagonal(diag)
+    end
+
+end #module
diff --git a/ext/DensityRatioEstimationJuMPExt.jl b/ext/DensityRatioEstimationJuMPExt.jl
@@ -0,0 +1,28 @@
+# ------------------------------------------------------------------
+# Licensed under the MIT License. See LICENSE in the project root.
+# ------------------------------------------------------------------
+module DensityRatioEstimationJuMPExt
+
+    if isdefined(Base, :get_extension)
+        using DensityRatioEstimation
+        using DensityRatioEstimation: LSIF, JuMPLib, AbstractKMM, uKMM, KMM
+        using DensityRatioEstimation.Parameters
+        using JuMP
+        using Ipopt
+        using LinearAlgebra
+        using Statistics
+    else
+        using ..DensityRatioEstimation
+        using ..DensityRatioEstimation: LSIF, JuMPLib, AbstractKMM, uKMM, KMM
+        using ..DensityRatioEstimation.Parameters
+        using ..JuMP
+        using ..Ipopt
+        using ..LinearAlgebra
+        using ..Statistics
+    end
+
+
+    include("../src/kmm/jump.jl")
+    include("../src/lsif/jump.jl")
+
+end #module
diff --git a/ext/DensityRatioEstimationOptimExt.jl b/ext/DensityRatioEstimationOptimExt.jl
@@ -0,0 +1,19 @@
+# ------------------------------------------------------------------
+# Licensed under the MIT License. See LICENSE in the project root.
+# ------------------------------------------------------------------
+module DensityRatioEstimationOptimExt
+
+    if isdefined(Base, :get_extension)
+        using DensityRatioEstimation
+        using DensityRatioEstimation: KLIEP, LSIF, OptimLib
+        using Optim
+    else
+        using ..DensityRatioEstimation
+        using ..DensityRatioEstimation: KLIEP, LSIF, OptimLib
+        using ..Optim
+    end
+    using LinearAlgebra
+    include("../src/kliep/optim.jl")
+    include("../src/lsif/optim.jl")
+
+end #module
diff --git a/src/DensityRatioEstimation.jl b/src/DensityRatioEstimation.jl
@@ -32,29 +32,43 @@ include("lcv.jl")
 # pure Julia implementations
 include("kmm/julia.jl")
 
-# implementations that require extra dependencies
-using Requires
-function __init__()
-  # KMM
-  @require JuMP = "4076af6c-e467-56ae-b986-b466b2749572" begin
-    @require Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9" include("kmm/jump.jl")
-  end
-
-  # KLIEP
-  @require Optim = "429524aa-4258-5aef-a3af-852621145aeb" include("kliep/optim.jl")
-  @require Convex = "f65535da-76fb-5f13-bab9-19810c17039a" begin
-    @require ECOS = "e2685f51-7e38-5353-a97d-a921fd2c8199" include("kliep/convex.jl")
-  end
 
-  # LSIF
-  @require Optim = "429524aa-4258-5aef-a3af-852621145aeb" include("lsif/optim.jl")
-  @require JuMP = "4076af6c-e467-56ae-b986-b466b2749572" begin
-    @require Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9" include("lsif/jump.jl")
+if !isdefined(Base,:get_extension)
+  using Requires
+end
+# implementations that require extra dependencies
+@static if !isdefined(Base,:get_extension)
+  function __init__()
+
+    #Solvers
+
+    #JuMP: KMM, LSIF
+    @require JuMP = "4076af6c-e467-56ae-b986-b466b2749572" begin
+      @require Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9" begin
+        include("../ext/DensityRatioEstimationJuMPExt.jl")
+      end 
+    end
+    #Optim: KLIEP, LSIF
+    @require Optim = "429524aa-4258-5aef-a3af-852621145aeb" begin
+      include("../ext/DensityRatioEstimationOptimExt.jl")
+    end
+
+    #Convex: KLIEP
+    @require Convex = "f65535da-76fb-5f13-bab9-19810c17039a" begin
+      @require ECOS = "e2685f51-7e38-5353-a97d-a921fd2c8199" begin
+        include("../ext/DensityRatioEstimationConvexExt.jl")
+      end
+    end
+
+    # AD and GPU libs
+    @require ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" begin 
+      include("../ext/DensityRatioEstimationChainRulesCoreExt.jl")
+    end
+
+    @require GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7" begin 
+      include("../ext/DensityRatioEstimationGPUArraysExt.jl")
+    end
   end
-
-  # AD and GPU libs
-  @require Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" include("utils/zygote.jl")
-  @require CuArrays = "3a865a2d-5b23-5a0f-bc46-62713ec82fae" include("utils/cuarrays.jl")
 end
 
 export

diff --git a/src/kliep/convex.jl b/src/kliep/convex.jl
@@ -2,10 +2,9 @@
 # Licensed under the MIT License. See LICENSE in the project root.
 # ------------------------------------------------------------------
 
-using .Convex
-using .ECOS
+#This file is part of the module DensityRatioEstimationConvexExt.
 
-function _kliep_coeffs(K_nu, K_de, dre::KLIEP, optlib::Type{ConvexLib})
+function DensityRatioEstimation._kliep_coeffs(K_nu, K_de, dre::KLIEP, optlib::Type{ConvexLib})
   # retrieve parameters
   σ, b = dre.σ, size(K_de, 2)
 

diff --git a/src/kliep/optim.jl b/src/kliep/optim.jl
@@ -2,9 +2,9 @@
 # Licensed under the MIT License. See LICENSE in the project root.
 # ------------------------------------------------------------------
 
-using .Optim
+#This file is part of the module DensityRatioEstimationOptimExt.
 
-function _kliep_coeffs(K_nu, K_de, dre::KLIEP, optlib::Type{OptimLib})
+function DensityRatioEstimation._kliep_coeffs(K_nu, K_de, dre::KLIEP, optlib::Type{OptimLib})
   # retrieve parameters
   σ, b = dre.σ, size(K_de, 2)
 

diff --git a/src/kmm/jump.jl b/src/kmm/jump.jl
@@ -2,8 +2,7 @@
 # Licensed under the MIT License. See LICENSE in the project root.
 # ------------------------------------------------------------------
 
-using .JuMP
-using .Ipopt
+#This file is part of the module DensityRatioEstimationJuMPExt.
 
 function _kmm_jump_model(K, κ, dre::AbstractKMM, optlib::Type{JuMPLib})
   # number of denominator samples
@@ -17,7 +16,7 @@ function _kmm_jump_model(K, κ, dre::AbstractKMM, optlib::Type{JuMPLib})
   return model, β
 end
 
-function _kmm_ratios(K, κ, dre::uKMM, optlib::Type{JuMPLib})
+function DensityRatioEstimation._kmm_ratios(K, κ, dre::uKMM, optlib::Type{JuMPLib})
   # build the problem without constraints
   model, β = _kmm_jump_model(K, κ, dre, optlib)
 
@@ -28,7 +27,7 @@ function _kmm_ratios(K, κ, dre::uKMM, optlib::Type{JuMPLib})
   value.(β)
 end
 
-function _kmm_ratios(K, κ, dre::KMM, optlib::Type{JuMPLib})
+function DensityRatioEstimation._kmm_ratios(K, κ, dre::KMM, optlib::Type{JuMPLib})
   # retrieve parameters
   @unpack B, ϵ = dre
 

diff --git a/src/lsif/jump.jl b/src/lsif/jump.jl
@@ -2,10 +2,9 @@
 # Licensed under the MIT License. See LICENSE in the project root.
 # ------------------------------------------------------------------
 
-using .JuMP
-using .Ipopt
+#This file is part of the module DensityRatioEstimationJuMPExt.
 
-function _lsif_coeffs(H, h, dre::LSIF, optlib::Type{JuMPLib})
+function DensityRatioEstimation._lsif_coeffs(H, h, dre::LSIF, optlib::Type{JuMPLib})
   # retrieve parameters
   λ, b = dre.λ, length(h)
 

diff --git a/src/lsif/optim.jl b/src/lsif/optim.jl
@@ -2,9 +2,9 @@
 # Licensed under the MIT License. See LICENSE in the project root.
 # ------------------------------------------------------------------
 
-using .Optim
+#This file is part of the module DensityRatioEstimationOptimExt.
 
-function _lsif_coeffs(H, h, dre::LSIF, optlib::Type{OptimLib})
+function DensityRatioEstimation._lsif_coeffs(H, h, dre::LSIF, optlib::Type{OptimLib})
   # retrieve parameters
   λ, b = dre.λ, length(h)
 

diff --git a/src/utils/cuarrays.jl b/src/utils/cuarrays.jl
diff --git a/src/utils/zygote.jl b/src/utils/zygote.jl