[dask] Fixed Dask type annotation (#4558)

python-package/lightgbm/dask.py

@@ -396,10 +396,10 @@ def _train(
     group: Optional[_DaskVectorLike] = None,
     eval_set: Optional[List[Tuple[_DaskMatrixLike, _DaskCollection]]] = None,
     eval_names: Optional[List[str]] = None,
-    eval_sample_weight: Optional[List[_DaskCollection]] = None,
+    eval_sample_weight: Optional[List[_DaskVectorLike]] = None,
     eval_class_weight: Optional[List[Union[dict, str]]] = None,
-    eval_init_score: Optional[List[_DaskCollection]] = None,
-    eval_group: Optional[List[_DaskCollection]] = None,
+    eval_init_score: Optional[List[_DaskVectorLike]] = None,
+    eval_group: Optional[List[_DaskVectorLike]] = None,
     eval_metric: Optional[Union[Callable, str, List[Union[Callable, str]]]] = None,
     eval_at: Optional[Iterable[int]] = None,
     **kwargs: Any
@@ -435,13 +435,13 @@ def _train(
         of evals_result_ and best_score_ will be 'not_evaluated'.
     eval_names : list of str, or None, optional (default=None)
         Names of eval_set.
-    eval_sample_weight : list of Dask Arrays, Dask Series or None, optional (default=None)
+    eval_sample_weight : list of Dask Arrays or Dask Series, or None, optional (default=None)
         Weights for each validation set in eval_set.
     eval_class_weight : list of dict or str, or None, optional (default=None)
         Class weights, one dict or str for each validation set in eval_set.
-    eval_init_score : list of Dask Arrays, Dask Series or None, optional (default=None)
+    eval_init_score : list of Dask Arrays or Dask Series, or None, optional (default=None)
         Initial model score for each validation set in eval_set.
-    eval_group : list of Dask Arrays, Dask Series or None, optional (default=None)
+    eval_group : list of Dask Arrays or Dask Series, or None, optional (default=None)
         Group/query for each validation set in eval_set.
     eval_metric : str, callable, list or None, optional (default=None)
         If str, it should be a built-in evaluation metric to use.
@@ -1025,10 +1025,10 @@ class _DaskLGBMModel:
         group: Optional[_DaskVectorLike] = None,
         eval_set: Optional[List[Tuple[_DaskMatrixLike, _DaskCollection]]] = None,
         eval_names: Optional[List[str]] = None,
-        eval_sample_weight: Optional[List[_DaskCollection]] = None,
+        eval_sample_weight: Optional[List[_DaskVectorLike]] = None,
         eval_class_weight: Optional[List[Union[dict, str]]] = None,
-        eval_init_score: Optional[List[_DaskCollection]] = None,
-        eval_group: Optional[List[_DaskCollection]] = None,
+        eval_init_score: Optional[List[_DaskVectorLike]] = None,
+        eval_group: Optional[List[_DaskVectorLike]] = None,
         eval_metric: Optional[Union[Callable, str, List[Union[Callable, str]]]] = None,
         eval_at: Optional[Iterable[int]] = None,
         early_stopping_rounds: Optional[int] = None,
@@ -1162,9 +1162,9 @@ class DaskLGBMClassifier(LGBMClassifier, _DaskLGBMModel):
         init_score: Optional[_DaskVectorLike] = None,
         eval_set: Optional[List[Tuple[_DaskMatrixLike, _DaskCollection]]] = None,
         eval_names: Optional[List[str]] = None,
-        eval_sample_weight: Optional[List[_DaskCollection]] = None,
+        eval_sample_weight: Optional[List[_DaskVectorLike]] = None,
         eval_class_weight: Optional[List[Union[dict, str]]] = None,
-        eval_init_score: Optional[List[_DaskCollection]] = None,
+        eval_init_score: Optional[List[_DaskVectorLike]] = None,
         eval_metric: Optional[Union[Callable, str, List[Union[Callable, str]]]] = None,
         early_stopping_rounds: Optional[int] = None,
         **kwargs: Any
@@ -1194,9 +1194,9 @@ class DaskLGBMClassifier(LGBMClassifier, _DaskLGBMModel):
         sample_weight_shape="Dask Array or Dask Series of shape = [n_samples] or None, optional (default=None)",
         init_score_shape="Dask Array or Dask Series of shape = [n_samples] or None, optional (default=None)",
         group_shape="Dask Array or Dask Series or None, optional (default=None)",
-        eval_sample_weight_shape="list of Dask Arrays or Dask Series or None, optional (default=None)",
-        eval_init_score_shape="list of Dask Arrays or Dask Series or None, optional (default=None)",
-        eval_group_shape="list of Dask Arrays or Dask Series or None, optional (default=None)"
+        eval_sample_weight_shape="list of Dask Arrays or Dask Series, or None, optional (default=None)",
+        eval_init_score_shape="list of Dask Arrays or Dask Series, or None, optional (default=None)",
+        eval_group_shape="list of Dask Arrays or Dask Series, or None, optional (default=None)"
     )
 
     # DaskLGBMClassifier does not support group, eval_group, early_stopping_rounds.
@@ -1341,8 +1341,8 @@ class DaskLGBMRegressor(LGBMRegressor, _DaskLGBMModel):
         init_score: Optional[_DaskVectorLike] = None,
         eval_set: Optional[List[Tuple[_DaskMatrixLike, _DaskCollection]]] = None,
         eval_names: Optional[List[str]] = None,
-        eval_sample_weight: Optional[List[_DaskCollection]] = None,
-        eval_init_score: Optional[List[_DaskCollection]] = None,
+        eval_sample_weight: Optional[List[_DaskVectorLike]] = None,
+        eval_init_score: Optional[List[_DaskVectorLike]] = None,
         eval_metric: Optional[Union[Callable, str, List[Union[Callable, str]]]] = None,
         early_stopping_rounds: Optional[int] = None,
         **kwargs: Any
@@ -1371,9 +1371,9 @@ class DaskLGBMRegressor(LGBMRegressor, _DaskLGBMModel):
         sample_weight_shape="Dask Array or Dask Series of shape = [n_samples] or None, optional (default=None)",
         init_score_shape="Dask Array or Dask Series of shape = [n_samples] or None, optional (default=None)",
         group_shape="Dask Array or Dask Series or None, optional (default=None)",
-        eval_sample_weight_shape="list of Dask Arrays or Dask Series or None, optional (default=None)",
-        eval_init_score_shape="list of Dask Arrays or Dask Series or None, optional (default=None)",
-        eval_group_shape="list of Dask Arrays or Dask Series or None, optional (default=None)"
+        eval_sample_weight_shape="list of Dask Arrays or Dask Series, or None, optional (default=None)",
+        eval_init_score_shape="list of Dask Arrays or Dask Series, or None, optional (default=None)",
+        eval_group_shape="list of Dask Arrays or Dask Series, or None, optional (default=None)"
     )
 
     # DaskLGBMRegressor does not support group, eval_class_weight, eval_group, early_stopping_rounds.
@@ -1503,9 +1503,9 @@ class DaskLGBMRanker(LGBMRanker, _DaskLGBMModel):
         group: Optional[_DaskVectorLike] = None,
         eval_set: Optional[List[Tuple[_DaskMatrixLike, _DaskCollection]]] = None,
         eval_names: Optional[List[str]] = None,
-        eval_sample_weight: Optional[List[_DaskCollection]] = None,
-        eval_init_score: Optional[List[_DaskCollection]] = None,
-        eval_group: Optional[List[_DaskCollection]] = None,
+        eval_sample_weight: Optional[List[_DaskVectorLike]] = None,
+        eval_init_score: Optional[List[_DaskVectorLike]] = None,
+        eval_group: Optional[List[_DaskVectorLike]] = None,
         eval_metric: Optional[Union[Callable, str, List[Union[Callable, str]]]] = None,
         eval_at: Iterable[int] = (1, 2, 3, 4, 5),
         early_stopping_rounds: Optional[int] = None,
@@ -1538,9 +1538,9 @@ class DaskLGBMRanker(LGBMRanker, _DaskLGBMModel):
         sample_weight_shape="Dask Array or Dask Series of shape = [n_samples] or None, optional (default=None)",
         init_score_shape="Dask Array or Dask Series of shape = [n_samples] or None, optional (default=None)",
         group_shape="Dask Array or Dask Series or None, optional (default=None)",
-        eval_sample_weight_shape="list of Dask Arrays or Dask Series or None, optional (default=None)",
-        eval_init_score_shape="list of Dask Arrays or Dask Series or None, optional (default=None)",
-        eval_group_shape="list of Dask Arrays or Dask Series or None, optional (default=None)"
+        eval_sample_weight_shape="list of Dask Arrays or Dask Series, or None, optional (default=None)",
+        eval_init_score_shape="list of Dask Arrays or Dask Series, or None, optional (default=None)",
+        eval_group_shape="list of Dask Arrays or Dask Series, or None, optional (default=None)"
     )
 
     # DaskLGBMRanker does not support eval_class_weight or early stopping