diff --git a/Project.toml b/Project.toml
index 099648e..8314503 100644
--- a/Project.toml
+++ b/Project.toml
@@ -4,6 +4,8 @@ authors = ["Kai Xu <xukai921110@gmail.com>", "Júlio Hoffimann <julio.hoffimann@
 version = "1.2.5"
 
 [deps]
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Parameters = "d96e819e-fc66-5662-9728-84c9c7592b0a"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -11,26 +13,22 @@ 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"
 DensityRatioEstimationConvexExt = ["Convex", "ECOS"]
-DensityRatioEstimationGPUArraysExt = "GPUArrays"
 DensityRatioEstimationJuMPExt = ["JuMP", "Ipopt"]
 DensityRatioEstimationOptimExt = "Optim"
 
 [compat]
-ChainRulesCore = "1"
+ChainRulesCore = "1.25.0"
 Convex = "0.15, 0.16"
 ECOS = "1"
-GPUArrays = "8, 9, 10, 11"
+GPUArraysCore = "0.2.0"
 Ipopt = "1"
 JuMP = "1"
 LinearAlgebra = "1.9"
@@ -40,12 +38,3 @@ Random = "1.9"
 Statistics = "1.9"
 StatsBase = "0.33, 0.34"
 julia = "1.9"
-
-[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
deleted file mode 100644
index a7949e0..0000000
--- a/ext/DensityRatioEstimationChainRulesCoreExt.jl
+++ /dev/null
@@ -1,12 +0,0 @@
-# ------------------------------------------------------------------
-# Licensed under the MIT License. See LICENSE in the project root.
-# ------------------------------------------------------------------
-
-module DensityRatioEstimationChainRulesCoreExt
-
-using DensityRatioEstimation
-using ChainRulesCore
-
-ChainRulesCore.@non_differentiable DensityRatioEstimation.safe_diagm(::Any, ::Any)
-
-end #module
diff --git a/ext/DensityRatioEstimationConvexExt.jl b/ext/DensityRatioEstimationConvexExt/DensityRatioEstimationConvexExt.jl
similarity index 84%
rename from ext/DensityRatioEstimationConvexExt.jl
rename to ext/DensityRatioEstimationConvexExt/DensityRatioEstimationConvexExt.jl
index df177c7..c6e716f 100644
--- a/ext/DensityRatioEstimationConvexExt.jl
+++ b/ext/DensityRatioEstimationConvexExt/DensityRatioEstimationConvexExt.jl
@@ -4,11 +4,12 @@
 
 module DensityRatioEstimationConvexExt
 
-using DensityRatioEstimation
 using DensityRatioEstimation: KLIEP, ConvexLib
 using Convex
 using ECOS
 
-include("../src/kliep/convex.jl")
+import DensityRatioEstimation
+
+include("kliep.jl")
 
 end #module
diff --git a/src/kliep/convex.jl b/ext/DensityRatioEstimationConvexExt/kliep.jl
similarity index 93%
rename from src/kliep/convex.jl
rename to ext/DensityRatioEstimationConvexExt/kliep.jl
index 1fedb95..b5bbb90 100644
--- a/src/kliep/convex.jl
+++ b/ext/DensityRatioEstimationConvexExt/kliep.jl
@@ -2,8 +2,6 @@
 # Licensed under the MIT License. See LICENSE in the project root.
 # ------------------------------------------------------------------
 
-# This file is part of the module DensityRatioEstimationConvexExt.
-
 function DensityRatioEstimation._kliep_coeffs(K_nu, K_de, dre::KLIEP, optlib::Type{ConvexLib})
   # retrieve parameters
   σ, b = dre.σ, size(K_de, 2)
diff --git a/ext/DensityRatioEstimationGPUArraysExt.jl b/ext/DensityRatioEstimationGPUArraysExt.jl
deleted file mode 100644
index 4b5ab02..0000000
--- a/ext/DensityRatioEstimationGPUArraysExt.jl
+++ /dev/null
@@ -1,18 +0,0 @@
-# ------------------------------------------------------------------
-# Licensed under the MIT License. See LICENSE in the project root.
-# ------------------------------------------------------------------
-
-module DensityRatioEstimationGPUArraysExt
-
-using DensityRatioEstimation
-using GPUArrays
-using LinearAlgebra
-
-# Avoid `mat + a * I` with CUDA which involves 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/DensityRatioEstimationJuMPExt.jl
similarity index 82%
rename from ext/DensityRatioEstimationJuMPExt.jl
rename to ext/DensityRatioEstimationJuMPExt/DensityRatioEstimationJuMPExt.jl
index 4626dfc..7a467af 100644
--- a/ext/DensityRatioEstimationJuMPExt.jl
+++ b/ext/DensityRatioEstimationJuMPExt/DensityRatioEstimationJuMPExt.jl
@@ -4,7 +4,6 @@
 
 module DensityRatioEstimationJuMPExt
 
-using DensityRatioEstimation
 using DensityRatioEstimation: LSIF, JuMPLib, AbstractKMM, uKMM, KMM
 using DensityRatioEstimation.Parameters
 using JuMP
@@ -12,7 +11,9 @@ using Ipopt
 using LinearAlgebra
 using Statistics
 
-include("../src/kmm/jump.jl")
-include("../src/lsif/jump.jl")
+import DensityRatioEstimation
+
+include("kmm.jl")
+include("lsif.jl")
 
 end #module
diff --git a/src/kmm/jump.jl b/ext/DensityRatioEstimationJuMPExt/kmm.jl
similarity index 100%
rename from src/kmm/jump.jl
rename to ext/DensityRatioEstimationJuMPExt/kmm.jl
diff --git a/src/lsif/jump.jl b/ext/DensityRatioEstimationJuMPExt/lsif.jl
similarity index 100%
rename from src/lsif/jump.jl
rename to ext/DensityRatioEstimationJuMPExt/lsif.jl
diff --git a/ext/DensityRatioEstimationOptimExt.jl b/ext/DensityRatioEstimationOptimExt/DensityRatioEstimationOptimExt.jl
similarity index 78%
rename from ext/DensityRatioEstimationOptimExt.jl
rename to ext/DensityRatioEstimationOptimExt/DensityRatioEstimationOptimExt.jl
index 77578fb..d2ca642 100644
--- a/ext/DensityRatioEstimationOptimExt.jl
+++ b/ext/DensityRatioEstimationOptimExt/DensityRatioEstimationOptimExt.jl
@@ -4,13 +4,13 @@
 
 module DensityRatioEstimationOptimExt
 
-using DensityRatioEstimation
 using DensityRatioEstimation: KLIEP, LSIF, OptimLib
+using LinearAlgebra
 using Optim
 
-using LinearAlgebra
+import DensityRatioEstimation
 
-include("../src/kliep/optim.jl")
-include("../src/lsif/optim.jl")
+include("kliep.jl")
+include("lsif.jl")
 
 end #module
diff --git a/src/kliep/optim.jl b/ext/DensityRatioEstimationOptimExt/kliep.jl
similarity index 100%
rename from src/kliep/optim.jl
rename to ext/DensityRatioEstimationOptimExt/kliep.jl
diff --git a/src/lsif/optim.jl b/ext/DensityRatioEstimationOptimExt/lsif.jl
similarity index 100%
rename from src/lsif/optim.jl
rename to ext/DensityRatioEstimationOptimExt/lsif.jl
diff --git a/src/DensityRatioEstimation.jl b/src/DensityRatioEstimation.jl
index 6f9e24b..5a81df2 100644
--- a/src/DensityRatioEstimation.jl
+++ b/src/DensityRatioEstimation.jl
@@ -10,7 +10,9 @@ using LinearAlgebra
 using Parameters
 using Random
 
-# implement fit for estimators
+using GPUArraysCore: AbstractGPUMatrix
+using ChainRulesCore: @non_differentiable
+
 import StatsBase: fit
 
 # API for density ratio estimation
@@ -29,9 +31,6 @@ include("lsif.jl")
 # available estimator fitters
 include("lcv.jl")
 
-# pure Julia implementations
-include("kmm/julia.jl")
-
 export
   # optim libs
   OptimizationLibrary,
diff --git a/src/kmm.jl b/src/kmm.jl
index 84da514..b208ae8 100644
--- a/src/kmm.jl
+++ b/src/kmm.jl
@@ -88,3 +88,9 @@ default_optlib(dre::Type{<:KMM}) = JuMPLib
 available_optlib(dre::Type{<:KMM}) = [JuMPLib]
 
 _kmm_ratios(K, κ, dre::AbstractKMM, optlib::Type{<:OptimizationLibrary}) = _throw_opt_error(dre, optlib)
+
+# pure Julia implementation
+function _kmm_ratios(K, κ, dre::uKMM, optlib::Type{JuliaLib})
+  # density ratio via solver
+  K \ vec(κ)
+end
diff --git a/src/kmm/julia.jl b/src/kmm/julia.jl
deleted file mode 100644
index 1dd3d37..0000000
--- a/src/kmm/julia.jl
+++ /dev/null
@@ -1,8 +0,0 @@
-# ------------------------------------------------------------------
-# Licensed under the MIT License. See LICENSE in the project root.
-# ------------------------------------------------------------------
-
-function _kmm_ratios(K, κ, dre::uKMM, optlib::Type{JuliaLib})
-  # density ratio via solver
-  K \ vec(κ)
-end
diff --git a/src/utils.jl b/src/utils.jl
index f6aad8b..fbdbc8b 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -9,8 +9,6 @@ Return the Euclidean distance between two indexable objects.
 """
 euclidsq(x, y) = sum((x[i] - y[i])^2 for i in eachindex(x))
 
-# Support matrix data in a GPU and AD compatible way
-
 """
     euclidsq(X::T, Y::T) where {T<:AbstractMatrix}
 
@@ -55,13 +53,18 @@ gaussian_gramian(esq, σ::AbstractFloat) = exp.(-esq ./ 2σ^2)
 Generate a squared matrix whose diagonal is `a` that is 
 compatible to perform addition on `mat`. It behaves 
 differently based on whether `mat` is on a CPU or GPU.
-
-It is compatible with
-- CuArrays.jl (see lib/cuarrays.jl)
-- Zygote.jl (see lib/zygote.jl)
 """
 safe_diagm(mat, a) = a * I
 
+# avoid `mat + a * I` on GPU which involves scalar operations and is slow
+function safe_diagm(mat::AbstractGPUMatrix, a)
+  diag = similar(mat, size(m, 1))
+  fill!(diag, a)
+  Diagonal(diag)
+end
+
+@non_differentiable safe_diagm(::Any, ::Any)
+
 ###################################################
 ##   Functions and objects for throwing errors   ##
 ###################################################