File: test_large_input.py

package info (click to toggle)
xgboost 3.0.0-1
  • links: PTS, VCS
  • area: main
  • in suites: trixie
  • size: 13,796 kB
  • sloc: cpp: 67,502; python: 35,503; java: 4,676; ansic: 1,426; sh: 1,320; xml: 1,197; makefile: 204; javascript: 19
file content (24 lines) | stat: -rw-r--r-- 670 bytes parent folder | download | duplicates (2)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
import cupy as cp
import numpy as np
import pytest

import xgboost as xgb


# Test for integer overflow or out of memory exceptions
def test_large_input():
    available_bytes, _ = cp.cuda.runtime.memGetInfo()
    # 15 GB
    required_bytes = 1.5e10
    if available_bytes < required_bytes:
        pytest.skip("Not enough memory on this device")
    n = 1000
    m = ((1 << 31) + n - 1) // n
    assert np.log2(m * n) > 31
    X = cp.ones((m, n), dtype=np.float32)
    y = cp.ones(m)
    w = cp.ones(m)
    dmat = xgb.QuantileDMatrix(X, y, weight=w)
    booster = xgb.train({"tree_method": "gpu_hist", "max_depth": 1}, dmat, 1)
    del y
    booster.inplace_predict(X)