scikit-learn
diff --git a/‎sklearn/feature_selection/_mutual_info.py
Copy file name to clipboardExpand all lines: sklearn/feature_selection/_mutual_info.py
+10-6 b/‎sklearn/feature_selection/_mutual_info.py
Copy file name to clipboardExpand all lines: sklearn/feature_selection/_mutual_info.py
+10-6
@@ -5,10 +5,11 @@
 
 import numpy as np
 from scipy.sparse import issparse
+from scipy.spatial import KDTree
 from scipy.special import digamma
 
 from ..metrics.cluster import mutual_info_score
-from ..neighbors import KDTree, NearestNeighbors
+from ..neighbors import NearestNeighbors
 from ..preprocessing import scale
 from ..utils import check_random_state
 from ..utils._param_validation import Interval, StrOptions, validate_params
@@ -62,12 +63,14 @@ def _compute_mi_cc(x, y, n_neighbors):
 
     # KDTree is explicitly fit to allow for the querying of number of
     # neighbors within a specified radius
-    kd = KDTree(x, metric="chebyshev")
-    nx = kd.query_radius(x, radius, count_only=True, return_distance=False)
+    kd = KDTree(x)
+    nx = kd.query_ball_point(x, radius, p=np.inf)
+    nx = [len(sub_list) for sub_list in nx]
     nx = np.array(nx) - 1.0
 
-    kd = KDTree(y, metric="chebyshev")
-    ny = kd.query_radius(y, radius, count_only=True, return_distance=False)
+    kd = KDTree(y)
+    ny = kd.query_ball_point(y, radius, p=np.inf)
+    ny = [len(sub_list) for sub_list in ny]
     ny = np.array(ny) - 1.0
 
     mi = (
@@ -140,7 +143,8 @@ def _compute_mi_cd(c, d, n_neighbors):
     radius = radius[mask]
 
     kd = KDTree(c)
-    m_all = kd.query_radius(c, radius, count_only=True, return_distance=False)
+    m_all = kd.query_ball_point(c, radius)
+    m_all = [len(sub_list) for sub_list in m_all]
     m_all = np.array(m_all)
 
     mi = (