From 2503c9bb77a03d5aa1356bbddf1938c139597e6d Mon Sep 17 00:00:00 2001
From: vic <viclafargue@nvidia.com>
Date: Tue, 11 Aug 2026 11:04:31 +0200
Subject: [PATCH] Fix sporadic KBinsDiscretizer mismatches with scikit-learn

---
 .../sklearn/preprocessing/_discretization.py  | 19 ++----
 python/cuml/tests/test_preprocessing.py       | 64 ++++++++++++++-----
 2 files changed, 52 insertions(+), 31 deletions(-)

diff --git a/python/cuml/cuml/_thirdparty/sklearn/preprocessing/_discretization.py b/python/cuml/cuml/_thirdparty/sklearn/preprocessing/_discretization.py
index 84cf41bd46..0e389876d8 100644
--- a/python/cuml/cuml/_thirdparty/sklearn/preprocessing/_discretization.py
+++ b/python/cuml/cuml/_thirdparty/sklearn/preprocessing/_discretization.py
@@ -31,10 +31,6 @@
 from ..utils.validation import FLOAT_DTYPES
 
 
-def digitize(x, bins):
-    return np.searchsorted(bins, x, side='left')
-
-
 class KBinsDiscretizer(TransformerMixin,
                        BaseEstimator,
                        SparseInputTagMixin):
@@ -198,7 +194,7 @@ def fit(self, X, y=None) -> "KBinsDiscretizer":
                 continue
 
             if self.strategy == 'uniform':
-                bin_edges[jj] = np.linspace(col_min, col_max, n_bins[jj] + 1)
+                bin_edges[jj] = np.linspace(col_min, col_max, int(n_bins[jj]) + 1)
 
             elif self.strategy == 'quantile':
                 quantiles = np.linspace(0, 100, n_bins[jj] + 1)
@@ -297,16 +293,9 @@ def transform(self, X):
 
         bin_edges = self.bin_edges_
         for jj in range(Xt.shape[1]):
-            # Values which are close to a bin edge are susceptible to numeric
-            # instability. Add eps to X so these values are binned correctly
-            # with respect to their decimal truncation. See documentation of
-            # numpy.isclose for an explanation of ``rtol`` and ``atol``.
-            rtol = 1.e-5
-            atol = 1.e-8
-            eps = atol + rtol * np.abs(Xt[:, jj])
-            Xt[:, jj] = digitize(Xt[:, jj] + eps, bin_edges[jj][1:])
-        self.n_bins_ = np.asarray(self.n_bins_)
-        np.clip(Xt, 0, self.n_bins_ - 1, out=Xt)
+            Xt[:, jj] = np.searchsorted(
+                bin_edges[jj][1:-1], Xt[:, jj], side='right'
+            )
 
         Xt = Xt.astype(np.int32)
         if self.encode == 'ordinal':
diff --git a/python/cuml/tests/test_preprocessing.py b/python/cuml/tests/test_preprocessing.py
index 632e791bd6..9fb4d39a9f 100644
--- a/python/cuml/tests/test_preprocessing.py
+++ b/python/cuml/tests/test_preprocessing.py
@@ -8,6 +8,7 @@
 import pytest
 import scipy
 import sklearn
+from packaging.version import Version
 from sklearn.impute import MissingIndicator as skMissingIndicator
 from sklearn.impute import SimpleImputer as skSimpleImputer
 from sklearn.preprocessing import Binarizer as skBinarizer
@@ -787,22 +788,8 @@ def test_robust_scale_sparse(
 @pytest.mark.parametrize(
     "strategy",
     [
-        pytest.param(
-            "uniform",
-            marks=pytest.mark.xfail(
-                strict=False,
-                reason="Intermittent mismatch with sklearn"
-                " (https://github.com/rapidsai/cuml/issues/3481)",
-            ),
-        ),
-        pytest.param(
-            "quantile",
-            marks=pytest.mark.xfail(
-                strict=False,
-                reason="Intermittent mismatch with sklearn"
-                " (https://github.com/rapidsai/cuml/issues/2933)",
-            ),
-        ),
+        "uniform",
+        "quantile",
         "kmeans",
     ],
 )
@@ -825,6 +812,14 @@ def test_kbinsdiscretizer(
         assert type(t_X) is type(X)
         assert type(r_X) is type(t_X)
 
+    sklearn_kwargs = {}
+    if strategy == "quantile" and Version(sklearn.__version__) >= Version(
+        "1.7"
+    ):
+        # cuML uses linear percentile interpolation. Scikit-learn exposed the
+        # method in 1.7 and changed its default in 1.9.
+        sklearn_kwargs["quantile_method"] = "linear"
+
     transformer = skKBinsDiscretizer(
         n_bins=n_bins,
         encode=encode,
@@ -833,6 +828,7 @@ def test_kbinsdiscretizer(
         subsample=200_000
         if strategy in ("uniform", "quantile", "kmeans")
         else None,
+        **sklearn_kwargs,
     )
     sk_t_X = transformer.fit_transform(X_np)
     sk_r_X = transformer.inverse_transform(sk_t_X)
@@ -844,6 +840,42 @@ def test_kbinsdiscretizer(
         assert_allclose(r_X, sk_r_X)
 
 
+@pytest.mark.parametrize(
+    "X_fit,X_transform,n_bins",
+    [
+        pytest.param(
+            [0.0, 1.0],
+            [0.499999],
+            2,
+            id="value-below-bin-edge",
+        ),
+        # These float32 values put X_transform between edges produced when
+        # CuPy receives a NumPy integer or a Python integer for ``num``.
+        pytest.param(
+            [-11.525976, 9.953962],
+            [1.3619866],
+            20,
+            id="float32-linspace-num-promotion",
+        ),
+    ],
+)
+def test_kbinsdiscretizer_uniform_edge_parity(X_fit, X_transform, n_bins):
+    X_fit = np.asarray(X_fit, dtype=np.float32)[:, None]
+    X_transform = np.asarray(X_transform, dtype=np.float32)[:, None]
+
+    cu_transformer = cuKBinsDiscretizer(
+        n_bins=n_bins, encode="ordinal", strategy="uniform"
+    ).fit(cp.asarray(X_fit))
+    sk_transformer = skKBinsDiscretizer(
+        n_bins=n_bins, encode="ordinal", strategy="uniform"
+    ).fit(X_fit)
+
+    assert_allclose(
+        cu_transformer.transform(cp.asarray(X_transform)),
+        sk_transformer.transform(X_transform),
+    )
+
+
 @pytest.mark.parametrize("missing_values", [0, 1, np.nan])
 @pytest.mark.parametrize("features", ["missing-only", "all"])
 def test_missing_indicator(