Implement GPU predict leaf. (#6187)

tests/python/test_predict.py

@@ -23,6 +23,49 @@ def run_threaded_predict(X, rows, predict_func):
         assert f.result()
 
 
+def run_predict_leaf(predictor):
+    rows = 100
+    cols = 4
+    classes = 5
+    num_parallel_tree = 4
+    num_boost_round = 10
+    rng = np.random.RandomState(1994)
+    X = rng.randn(rows, cols)
+    y = rng.randint(low=0, high=classes, size=rows)
+    m = xgb.DMatrix(X, y)
+    booster = xgb.train(
+        {'num_parallel_tree': num_parallel_tree, 'num_class': classes,
+         'predictor': predictor, 'tree_method': 'hist'}, m,
+        num_boost_round=num_boost_round)
+
+    empty = xgb.DMatrix(np.ones(shape=(0, cols)))
+    empty_leaf = booster.predict(empty, pred_leaf=True)
+    assert empty_leaf.shape[0] == 0
+
+    leaf = booster.predict(m, pred_leaf=True)
+    assert leaf.shape[0] == rows
+    assert leaf.shape[1] == classes * num_parallel_tree * num_boost_round
+
+    for i in range(rows):
+        row = leaf[i, ...]
+        for j in range(num_boost_round):
+            start = classes * num_parallel_tree * j
+            end = classes * num_parallel_tree * (j + 1)
+            layer = row[start: end]
+            for c in range(classes):
+                tree_group = layer[c * num_parallel_tree:
+                                   (c+1) * num_parallel_tree]
+                assert tree_group.shape[0] == num_parallel_tree
+                # no subsampling so tree in same forest should output same
+                # leaf.
+                assert np.all(tree_group == tree_group[0])
+    return leaf
+
+
+def test_predict_leaf():
+    run_predict_leaf('cpu_predictor')
+
+
 class TestInplacePredict(unittest.TestCase):
     '''Tests for running inplace prediction'''
     def test_predict(self):