Merge branch 'master' of https://github.com/rudolfwilliam/DiffKDE

Author: rudolfwilliam

README.md

@@ -77,6 +77,7 @@ The current implementation provides the following functionality:
 | Kernels                  | Gaussian, Epanechnikov, Exponential, Tophat Approximation, von Mises-Fisher (data must lie on the unit sphere) |
 | Tree Algorithms          | Standard                    |
 | Bandwidths               | Float (Isotropic bandwidth matrix), Scott, Silverman |
+| Devices | CPU, GPU |
 
 </div>
 

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "torch-kde"
-version = "0.1.3"
+version = "0.1.4"
 description = "A differentiable implementation of kernel density estimation in PyTorch"
 readme = "README_PyPI.md"
 license = { text = "MIT" }

tests/test_kde.py

@@ -20,6 +20,8 @@
 TOLERANCE = 1e-1
 WEIGHTS = [False, True]
 
+DEVICES = ["cpu"]
+
 N1 = 100
 N2 = 10
 GRID_N = 1000

torchkde/modules.py

@@ -66,6 +66,7 @@ def __init__(
         self.algorithm = algorithm
         self.is_fitted = False
         self.n_samples = None
+        self.device = None
         self.n_features = None
         self.data = None
 
@@ -106,14 +107,17 @@ def fit(self,
         self.tree_.build(X)
         self.bandwidth = compute_bandwidth(X, self.bandwidth)
         self.kernel_module.bandwidth = self.bandwidth
+        self.device = X.device
         if sample_weight is None:
-            self.sample_weight = torch.ones(self.n_samples)
+            self.sample_weight = torch.ones(self.n_samples).to(self.device)
         else:
             assert sample_weight.shape[0] == self.n_samples, "Sample weights must have the same length as the data."
+            assert sample_weight.device == self.device, "Sample weights must be on the same device as the data."
             assert sample_weight.dim() == 1, "Sample weights must be one-dimensional."
             assert sample_weight.min() >= 0, "Sample weights must be non-negative."
             self.sample_weight = sample_weight
         self.is_fitted = True
+        self.device = X.device
 
         return self
 
@@ -136,6 +140,7 @@ def score_samples(self, X: torch.Tensor, batch_size: int = 128) -> torch.Tensor:
             data.
         """
         assert self.is_fitted, "Model must be fitted before scoring samples."
+        assert X.device == self.device, "Device of the query data must be on the same device as the data for fitting the estimator."
 
         n_samples = X.shape[0]
         # Compute log-density estimation with a kernel function