Test QDM with sparse data on CPU. (#9316)

This commit is contained in:
Jiaming Yuan 2023-06-19 21:27:03 +08:00 committed by GitHub
parent ee6809e642
commit 6d22ea793c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -253,9 +253,12 @@ class TestQuantileDMatrix:
self.run_ref_dmatrix(rng, "hist", True)
self.run_ref_dmatrix(rng, "hist", False)
def test_predict(self) -> None:
n_samples, n_features = 16, 2
X, y = make_categorical(n_samples, n_features, n_categories=13, onehot=False)
@pytest.mark.parametrize("sparsity", [0.0, 0.5])
def test_predict(self, sparsity: float) -> None:
n_samples, n_features = 256, 4
X, y = make_categorical(
n_samples, n_features, n_categories=13, onehot=False, sparsity=sparsity
)
Xy = xgb.DMatrix(X, y, enable_categorical=True)
booster = xgb.train({"tree_method": "hist"}, Xy)