Cleanup.

Author: trivialfis

python-package/xgboost/testing/__init__.py

@@ -511,19 +511,6 @@ def non_decreasing(L: Sequence[float], tolerance: float = 1e-4) -> bool:
     return all((y - x) >= -tolerance for x, y in zip(L, L[1:]))
 
 
-def predictor_equal(lhs: xgb.DMatrix, rhs: xgb.DMatrix) -> bool:
-    """Assert whether two DMatrices contain the same predictors."""
-    lcsr = lhs.get_data()
-    rcsr = rhs.get_data()
-    return all(
-        (
-            np.array_equal(lcsr.data, rcsr.data),
-            np.array_equal(lcsr.indices, rcsr.indices),
-            np.array_equal(lcsr.indptr, rcsr.indptr),
-        )
-    )
-
-
 M = TypeVar("M", xgb.Booster, xgb.XGBModel)
 
 

python-package/xgboost/testing/data_iter.py

@@ -8,6 +8,7 @@
 
 from ..compat import import_cupy
 from ..core import DataIter, DMatrix, ExtMemQuantileDMatrix, QuantileDMatrix
+from .utils import predictor_equal
 
 
 def run_mixed_sparsity(device: str) -> None:
@@ -36,7 +37,7 @@ def run_mixed_sparsity(device: str) -> None:
     y_arr = np.concatenate(y, axis=0)
     Xy_1 = QuantileDMatrix(X_arr, y_arr)
 
-    assert tm.predictor_equal(Xy_0, Xy_1)
+    assert predictor_equal(Xy_0, Xy_1)
 
 
 def check_invalid_cat_batches(device: str) -> None:

python-package/xgboost/testing/ordinal.py

@@ -23,7 +23,7 @@
     memory,
 )
 from .updater import get_basescore
-from .utils import Device, assert_allclose
+from .utils import Device, assert_allclose, predictor_equal
 
 
 @fcache
@@ -705,13 +705,7 @@ def run_recode_dmatrix(device: Device) -> None:
         assert cats_1 is not None
 
         assert cats_0.to_arrow() == cats_1.to_arrow()
-
-        csr_0 = Xy_0.get_data()
-        csr_1 = Xy_1.get_data()
-
-        np.testing.assert_allclose(csr_0.indptr, csr_1.indptr)
-        np.testing.assert_allclose(csr_0.indices, csr_1.indices)
-        np.testing.assert_allclose(csr_0.data, csr_1.data)
+        assert predictor_equal(Xy_0, Xy_1)
 
 
 def run_training_continuation(device: Device) -> None:

python-package/xgboost/testing/utils.py

@@ -5,6 +5,7 @@
 import numpy as np
 
 from ..compat import import_cupy
+from ..core import DMatrix
 from ..data import _is_cupy_alike
 
 Device: TypeAlias = Literal["cpu", "cuda"]
@@ -19,3 +20,16 @@ def assert_allclose(
     else:
         cp = import_cupy()
         cp.testing.assert_allclose(a, b, atol=atol, rtol=rtol)
+
+
+def predictor_equal(lhs: DMatrix, rhs: DMatrix) -> bool:
+    """Assert whether two DMatrices contain the same predictors."""
+    lcsr = lhs.get_data()
+    rcsr = rhs.get_data()
+    return all(
+        (
+            np.array_equal(lcsr.data, rcsr.data),
+            np.array_equal(lcsr.indices, rcsr.indices),
+            np.array_equal(lcsr.indptr, rcsr.indptr),
+        )
+    )

tests/python-gpu/test_device_quantile_dmatrix.py

@@ -12,6 +12,7 @@
     check_categorical_strings,
     check_ref_quantile_cut,
 )
+from xgboost.testing.utils import predictor_equal
 
 sys.path.append("tests/python")
 import test_quantile_dmatrix as tqd
@@ -70,7 +71,7 @@ def test_initialization(self, on_device: bool, device: str) -> None:
         )
         # query cuts from GIDX/Ellpack
         qXy = xgb.QuantileDMatrix(X[0], y[0], weight=w[0], max_bin=max_bin, ref=Xy)
-        tm.predictor_equal(Xy, qXy)
+        predictor_equal(Xy, qXy)
         with pytest.raises(ValueError, match="Inconsistent"):
             # max_bin changed.
             xgb.QuantileDMatrix(X[0], y[0], weight=w[0], max_bin=max_bin - 1, ref=Xy)
@@ -92,7 +93,7 @@ def test_initialization(self, on_device: bool, device: str) -> None:
         )
         # query cuts from GIDX/Ellpack
         qXy = xgb.QuantileDMatrix(X[0], y[0], weight=w[0], max_bin=max_bin, ref=Xy)
-        tm.predictor_equal(Xy, qXy)
+        predictor_equal(Xy, qXy)
         with pytest.raises(ValueError, match="Inconsistent"):
             # max_bin changed.
             xgb.QuantileDMatrix(X[0], y[0], weight=w[0], max_bin=max_bin - 1, ref=Xy)
@@ -247,7 +248,7 @@ def test_ltr(self) -> None:
         from_dm = xgb.QuantileDMatrix(X, weight=w, ref=Xy)
         from_qdm = xgb.QuantileDMatrix(X, weight=w, ref=Xy_qdm)
 
-        assert tm.predictor_equal(from_qdm, from_dm)
+        assert predictor_equal(from_qdm, from_dm)
 
     @pytest.mark.skipif(**tm.no_cupy())
     def test_check_inf(self) -> None:

tests/python/test_dmatrix.py

@@ -13,6 +13,7 @@
 from xgboost import testing as tm
 from xgboost.core import DataSplitMode
 from xgboost.testing.data import np_dtypes, run_base_margin_info
+from xgboost.testing.utils import predictor_equal
 
 dpath = "demo/data/"
 rng = np.random.RandomState(1994)
@@ -397,7 +398,7 @@ def test_dtypes(self) -> None:
         for orig, x in np_dtypes(n_samples, n_features):
             m0 = xgb.DMatrix(orig)
             m1 = xgb.DMatrix(x)
-            assert tm.predictor_equal(m0, m1)
+            assert predictor_equal(m0, m1)
 
 
 @pytest.mark.skipif(tm.is_windows(), reason="Rabit does not run on windows")

tests/python/test_quantile_dmatrix.py

@@ -13,14 +13,14 @@
     make_categorical,
     make_ltr,
     make_sparse_regression,
-    predictor_equal,
 )
 from xgboost.testing.data import check_inf, np_dtypes
 from xgboost.testing.data_iter import run_mixed_sparsity
 from xgboost.testing.quantile_dmatrix import (
     check_categorical_strings,
     check_ref_quantile_cut,
 )
+from xgboost.testing.utils import predictor_equal
 
 
 class TestQuantileDMatrix:

tests/python/test_with_pandas.py

@@ -7,6 +7,7 @@
 from xgboost import testing as tm
 from xgboost.core import DataSplitMode
 from xgboost.testing.data import pd_arrow_dtypes, pd_dtypes, run_base_margin_info
+from xgboost.testing.utils import predictor_equal
 
 try:
     import pandas as pd
@@ -482,9 +483,9 @@ def test_nullable_type(self, DMatrixT) -> None:
             if hasattr(orig.dtypes, "__iter__") and any(
                 dtype == "bool" for dtype in orig.dtypes
             ):
-                assert not tm.predictor_equal(m_orig, m_etype)
+                assert not predictor_equal(m_orig, m_etype)
             else:
-                assert tm.predictor_equal(m_orig, m_etype)
+                assert predictor_equal(m_orig, m_etype)
 
             np.testing.assert_allclose(m_orig.get_label(), m_etype.get_label())
             np.testing.assert_allclose(m_etype.get_label(), y.values.astype(np.float32))
@@ -511,7 +512,7 @@ def test_pyarrow_type(self, DMatrixT: Type[xgb.DMatrix]) -> None:
             m_orig = DMatrixT(orig, enable_categorical=True, label=y_orig)
             m_etype = DMatrixT(df, enable_categorical=True, label=y)
 
-            assert tm.predictor_equal(m_orig, m_etype)
+            assert predictor_equal(m_orig, m_etype)
             if y is not None:
                 np.testing.assert_allclose(m_orig.get_label(), m_etype.get_label())
                 np.testing.assert_allclose(m_etype.get_label(), y.values)

tests/python/test_with_scipy.py

@@ -7,7 +7,7 @@
 import scipy.sparse
 
 import xgboost as xgb
-from xgboost import testing as tm
+from xgboost.testing.utils import predictor_equal
 
 
 @pytest.mark.filterwarnings("error")
@@ -59,7 +59,7 @@ def test_csc(DMatrixT: Type[xgb.DMatrix], CSC: Type) -> None:
         data = np.array([0, 1, 2, 3, 4])
         row_idx = np.array([0, 1, 2, 0, 2])
         X = CSC((data, row_idx, indptr), shape=(3, 2))
-        assert tm.predictor_equal(DMatrixT(X.tocsr()), DMatrixT(X))
+        assert predictor_equal(DMatrixT(X.tocsr()), DMatrixT(X))
 
 
 @pytest.mark.filterwarnings("error")
@@ -84,4 +84,4 @@ def test_coo(DMatrixT: Type[xgb.DMatrix], COO: Type) -> None:
         assert dtrain.num_col() == 3
         assert dtrain.num_nonmissing() == data.size
 
-        assert tm.predictor_equal(DMatrixT(X.tocsr()), DMatrixT(X))
+        assert predictor_equal(DMatrixT(X.tocsr()), DMatrixT(X))