Hendrik/xgboost
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[backport] Ensure predict leaf output 1-dim vector where there's only 1 tree. (#6889) (#6937)

Browse Source
This commit is contained in:
Jiaming Yuan 2021-05-04 16:09:49 +08:00 committed by GitHub
parent 96f8843694
commit b78ad1e623
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 9 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
src/c_api/c_api_utils.h
View File

@ -98,6 +98,10 @@ inline void CalcPredictShape(bool strict_shape, PredictionType type, size_t rows
forest = std::max(static_cast<decltype(forest)>(1), forest);
shape[3] = forest;
*out_dim = shape.size();
} else if (chunksize == 1) {
*out_dim = 1;
shape.resize(*out_dim);
shape.front() = rows;
} else {
*out_dim = 2;
shape.resize(*out_dim);

5
tests/python/test_predict.py
View File

@ -75,6 +75,11 @@ def run_predict_leaf(predictor):
first = sliced[0, ...]
assert np.prod(first.shape) == classes * num_parallel_tree * ntree_limit
# When there's only 1 tree, the output is a 1 dim vector
booster = xgb.train({"tree_method": "hist"}, num_boost_round=1, dtrain=m)
assert booster.predict(m, pred_leaf=True).shape == (rows, )
return leaf