diff options
Diffstat (limited to 'src/python/gudhi/point_cloud/knn.py')
| -rw-r--r-- | src/python/gudhi/point_cloud/knn.py | 6 |
1 files changed, 5 insertions, 1 deletions
diff --git a/src/python/gudhi/point_cloud/knn.py b/src/python/gudhi/point_cloud/knn.py index 86008bc3..994be3b6 100644 --- a/src/python/gudhi/point_cloud/knn.py +++ b/src/python/gudhi/point_cloud/knn.py @@ -46,7 +46,7 @@ class KNearestNeighbors: sort_results (bool): if True, then distances and indices of each point are sorted on return, so that the first column contains the closest points. Otherwise, neighbors are returned in an arbitrary order. Defaults to True. - enable_autodiff (bool): if the input is a torch.tensor, jax.numpy.ndarray or tensorflow.Tensor, this + enable_autodiff (bool): if the input is a torch.tensor or tensorflow.Tensor, this instructs the function to compute distances in a way that works with automatic differentiation. This is experimental, not supported for all metrics, and requires the package EagerPy. Defaults to False. @@ -306,6 +306,10 @@ class KNearestNeighbors: if self.params["implementation"] == "ckdtree": qargs = {key: val for key, val in self.params.items() if key in {"p", "eps", "n_jobs"}} distances, neighbors = self.kdtree.query(X, k=self.k, **qargs) + if k == 1: + # SciPy decided to squeeze the last dimension for k=1 + distances = distances[:, None] + neighbors = neighbors[:, None] if self.return_index: if self.return_distance: return neighbors, distances |