diff --git a/python-package/xgboost/dask.py b/python-package/xgboost/dask.py index f203f88b27c8..500c96e50448 100644 --- a/python-package/xgboost/dask.py +++ b/python-package/xgboost/dask.py @@ -48,7 +48,6 @@ from .sklearn import _cls_predict_proba from .sklearn import XGBRanker - if TYPE_CHECKING: from dask import dataframe as dd from dask import array as da @@ -71,6 +70,20 @@ except ImportError: TrainReturnT = Dict[str, Any] # type:ignore +__all__ = [ + "RabitContext", + "DaskDMatrix", + "DaskDeviceQuantileDMatrix", + "DaskXGBRegressor", + "DaskXGBClassifier", + "DaskXGBRanker", + "DaskXGBRFRegressor", + "DaskXGBRFClassifier", + "train", + "predict", + "inplace_predict", +] + # TODOs: # - CV #