diff --git a/Project.toml b/Project.toml
index 2d0489c..fb985ec 100644
--- 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"
\ No newline at end of file
diff --git a/ext/DensityRatioEstimationChainRulesCoreExt.jl b/ext/DensityRatioEstimationChainRulesCoreExt.jl
new file mode 100644
index 0000000..70d63e7
--- /dev/null
+++ 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
new file mode 100644
index 0000000..97f9c3a
--- /dev/null
+++ 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
new file mode 100644
index 0000000..e3325f7
--- /dev/null
+++ 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
new file mode 100644
index 0000000..96b987d
--- /dev/null
+++ 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
new file mode 100644
index 0000000..7935313
--- /dev/null
+++ 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
index d9a27cb..447f82c 100644
--- 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
index ed97e72..d6f3bbf 100644
--- 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
index 26532d6..025fe63 100644
--- 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
index 855e201..13eb481 100644
--- 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
index aa1d100..3c4169e 100644
--- 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
index 9f5b8e7..9c7f27d 100644
--- 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
deleted file mode 100644
index c120dec..0000000
--- a/src/utils/cuarrays.jl
+++ /dev/null
@@ -1,8 +0,0 @@
-# ------------------------------------------------------------------
-# Licensed under the MIT License. See LICENSE in the project root.
-# ------------------------------------------------------------------
-
-import .CuArrays
-
-# Aovid `mat + a * I` with CuArrays which involes scalar operations and is slow
-safe_diagm(mat::CuArrays.CuArray, a::T) where {T} = CuArrays.CuArray{T}(diagm(0 => fill(a, size(mat, 1))))
diff --git a/src/utils/zygote.jl b/src/utils/zygote.jl
deleted file mode 100644
index 07fce47..0000000
--- a/src/utils/zygote.jl
+++ /dev/null
@@ -1,7 +0,0 @@
-# ------------------------------------------------------------------
-# Licensed under the MIT License. See LICENSE in the project root.
-# ------------------------------------------------------------------
-
-import .Zygote
-
-Zygote.@nograd safe_diagm