Make mypy strict

examples/custom_objective.py

@@ -3,7 +3,7 @@
 
 
 class MyMetric(lb.BaseMetric):
-    def metric(self, y: cn.ndarray, pred: cn.ndarray, w: cn.ndarray) -> float:
+    def metric(self, y: cn.ndarray, pred: cn.ndarray, w: cn.ndarray) -> cn.ndarray:
         return cn.sqrt(((y - pred) ** 2 * w).sum() / w.sum())
 
     def name(self):

legateboost/__init__.py

@@ -1,28 +1,5 @@
-from .legateboost import LBClassifier, LBRegressor
-from .metrics import (
-    BaseMetric,
-    ExponentialMetric,
-    GammaDevianceMetric,
-    LogLossMetric,
-    MSEMetric,
-    NormalCRPSMetric,
-    NormalLLMetric,
-    GammaLLMetric,
-    QuantileMetric,
-)
-from .objectives import (
-    BaseObjective,
-    ExponentialObjective,
-    GammaDevianceObjective,
-    GammaObjective,
-    LogLossObjective,
-    NormalObjective,
-    QuantileObjective,
-    SquaredErrorObjective,
-)
-from .callbacks import (
-    TrainingCallback,
-    EarlyStopping,
-)
-from .utils import mod_col_by_idx, pick_col_by_idx, set_col_by_idx
+from .legateboost import *
+from .metrics import *
+from .objectives import *
+from .callbacks import *
 from ._version import __version__

legateboost/callbacks.py

@@ -6,6 +6,8 @@
 from .legateboost import EvalResult, LBBase
 from .metrics import metrics
 
+__all__ = ["TrainingCallback", "EarlyStopping"]
+
 
 class TrainingCallback(ABC):
     """Interface for training callback."""
@@ -100,11 +102,11 @@ def _update_rounds(
     ) -> bool:
         def maximize(new: float, best: float) -> bool:
             """New score should be greater than the old one."""
-            return np.greater(new - self._min_delta, best)
+            return bool(np.greater(new - self._min_delta, best))
 
         def minimize(new: float, best: float) -> bool:
             """New score should be lesser than the old one."""
-            return np.greater(best - self._min_delta, new)
+            return bool(np.greater(best - self._min_delta, new))
 
         if metrics[metric].minimize():
             improve_op = minimize

legateboost/input_validation.py

@@ -1,10 +1,12 @@
-from typing import Any
+from typing import Any, List
 
 import numpy as np
 import scipy.sparse as sp
 
 import cunumeric as cn
 
+__all__: List[str] = []
+
 
 def check_sample_weight(sample_weight: Any, n: int) -> cn.ndarray:
     if sample_weight is None:

legateboost/legateboost.py

@@ -24,6 +24,8 @@
 
 EvalResult: TypeAlias = dict[str, dict[str, list[float]]]
 
+__all__ = ["LBBase", "LBClassifier", "LBRegressor"]
+
 
 class LBBase(BaseEstimator, PickleCunumericMixin):
     def __init__(

legateboost/library.py

@@ -66,7 +66,7 @@ def get_shared_library(self) -> str:
     def get_c_header(self) -> str:
         from .install_info import header
 
-        return header
+        return str(header)
 
     def get_registration_callback(self) -> str:
         return "legateboost_perform_registration"

legateboost/metrics.py

@@ -8,6 +8,18 @@
 from .special import erf, loggamma
 from .utils import pick_col_by_idx, sample_average, set_col_by_idx
 
+__all__ = [
+    "BaseMetric",
+    "MSEMetric",
+    "NormalLLMetric",
+    "GammaLLMetric",
+    "NormalCRPSMetric",
+    "GammaDevianceMetric",
+    "QuantileMetric",
+    "LogLossMetric",
+    "ExponentialMetric",
+]
+
 
 class BaseMetric(ABC):
     """The base class for metrics.
@@ -19,7 +31,7 @@ class BaseMetric(ABC):
     one = cn.ones(1, dtype=cn.float64)
 
     @abstractmethod
-    def metric(self, y: cn.ndarray, pred: cn.ndarray, w: cn.ndarray) -> float:
+    def metric(self, y: cn.ndarray, pred: cn.ndarray, w: cn.ndarray) -> cn.ndarray:
         """Computes the metric between the true labels `y` and predicted labels
         `pred`, weighted by `w`.
 
@@ -67,7 +79,7 @@ class MSEMetric(BaseMetric):
         :class:`legateboost.objectives.SquaredErrorObjective`
     """
 
-    def metric(self, y: cn.ndarray, pred: cn.ndarray, w: cn.ndarray) -> float:
+    def metric(self, y: cn.ndarray, pred: cn.ndarray, w: cn.ndarray) -> cn.ndarray:
         assert w.ndim == 1
         y = y.reshape(pred.shape)
         w_sum = w.sum()
@@ -109,7 +121,7 @@ class NormalLLMetric(BaseMetric):
         :class:`legateboost.objectives.NormalObjective`
     """  # noqa: E501
 
-    def metric(self, y: cn.ndarray, pred: cn.ndarray, w: cn.ndarray) -> float:
+    def metric(self, y: cn.ndarray, pred: cn.ndarray, w: cn.ndarray) -> cn.ndarray:
         y, pred = check_dist_param(y, pred)
         w_sum = w.sum()
         if w_sum == 0:
@@ -136,7 +148,7 @@ class GammaLLMetric(BaseMetric):
     predicted by the model."""
 
     @override
-    def metric(self, y: cn.ndarray, pred: cn.ndarray, w: cn.ndarray) -> float:
+    def metric(self, y: cn.ndarray, pred: cn.ndarray, w: cn.ndarray) -> cn.ndarray:
         y, pred = check_dist_param(y, pred)
 
         w_sum = w.sum()
@@ -174,15 +186,15 @@ class NormalCRPSMetric(BaseMetric):
         `Strictly Proper Scoring Rules, Prediction, and Estimation`
     """
 
-    def metric(self, y: cn.ndarray, pred: cn.ndarray, w: cn.ndarray) -> float:
+    def metric(self, y: cn.ndarray, pred: cn.ndarray, w: cn.ndarray) -> cn.ndarray:
         y, pred = check_dist_param(y, pred)
         loc = pred[:, :, 0]
         # `NormalObjective` outputs variance instead of scale.
         scale = cn.sqrt(pred[:, :, 1])
         z = (y - loc) / scale
         # This is negating the definition in [1] to make it a loss.
         v = scale * (z * (2 * norm_cdf(z) - 1) + 2 * norm_pdf(z) - 1 / cn.sqrt(cn.pi))
-        return sample_average(v, w)
+        return float(sample_average(v, w))
 
     def name(self) -> str:
         return "normal_crps"
@@ -194,7 +206,7 @@ class GammaDevianceMetric(BaseMetric):
     :math:`E = 2[(\\ln{\\frac{p_i}{y_i}} + \\frac{y_i}{p_i} - 1)w_i]`
     """
 
-    def metric(self, y: cn.ndarray, pred: cn.ndarray, w: cn.ndarray) -> float:
+    def metric(self, y: cn.ndarray, pred: cn.ndarray, w: cn.ndarray) -> cn.ndarray:
         eps = 1e-6
         y = y + eps
         pred = pred + eps
@@ -226,14 +238,14 @@ def __init__(self, quantiles: cn.ndarray = cn.array([0.25, 0.5, 0.75])) -> None:
         assert cn.all(0.0 <= quantiles) and cn.all(quantiles <= self.one)
         self.quantiles = quantiles
 
-    def metric(self, y: cn.ndarray, pred: cn.ndarray, w: cn.ndarray) -> float:
+    def metric(self, y: cn.ndarray, pred: cn.ndarray, w: cn.ndarray) -> cn.ndarray:
         assert w.ndim == 1
         assert y.shape[1] == 1
         assert pred.shape[1] == self.quantiles.size
         diff = y - pred
         indicator = diff <= 0
         loss = (self.quantiles[cn.newaxis, :] - indicator) * diff
-        return ((loss * w[:, cn.newaxis]).sum() / self.quantiles.size) / w.sum()
+        return float(((loss * w[:, cn.newaxis]).sum() / self.quantiles.size) / w.sum())
 
     def name(self) -> str:
         return "quantile"
@@ -258,7 +270,7 @@ class LogLossMetric(BaseMetric):
         :class:`legateboost.objectives.LogLossObjective`
     """  # noqa: E501
 
-    def metric(self, y: cn.ndarray, pred: cn.ndarray, w: cn.ndarray) -> float:
+    def metric(self, y: cn.ndarray, pred: cn.ndarray, w: cn.ndarray) -> cn.ndarray:
         y = y.squeeze()
         eps = cn.finfo(pred.dtype).eps
         cn.clip(pred, eps, 1 - eps, out=pred)
@@ -296,13 +308,13 @@ class ExponentialMetric(BaseMetric):
         :class:`legateboost.objectives.ExponentialObjective`
     """  # noqa: E501
 
-    def metric(self, y: cn.ndarray, pred: cn.ndarray, w: cn.ndarray) -> float:
+    def metric(self, y: cn.ndarray, pred: cn.ndarray, w: cn.ndarray) -> cn.ndarray:
         y = y.squeeze()
         # binary case
         if pred.ndim == 1 or pred.shape[1] == 1:
             pred = pred.squeeze()
             exp = cn.power(pred / (self.one - pred), 0.5 - y)
-            return float((exp * w).sum() / w.sum())
+            return (exp * w).sum() / w.sum()
 
         # multi-class case
         # note that exp loss is invariant to adding a constant to prediction
@@ -314,7 +326,7 @@ def metric(self, y: cn.ndarray, pred: cn.ndarray, w: cn.ndarray) -> float:
         # y_k[cn.arange(y.size), y.astype(cn.int32)] = 1.0
 
         exp = cn.exp(-1 / K * cn.sum(y_k * f, axis=1))
-        return float((exp * w).sum() / w.sum())
+        return (exp * w).sum() / w.sum()
 
     def name(self) -> str:
         return "exp"

legateboost/models/__init__.py

@@ -3,3 +3,6 @@
 from .krr import KRR
 from .nn import NN
 from .base_model import BaseModel
+from typing import List
+
+__all__: List[str] = ["Tree", "Linear", "KRR", "NN", "BaseModel"]

legateboost/models/krr.py

@@ -219,6 +219,7 @@ def __str__(self) -> str:
     def __eq__(self, other: object) -> bool:
         if not isinstance(other, KRR):
             raise NotImplementedError()
-        return (other.betas_ == self.betas_).all() and (
-            other.X_train == self.X_train
-        ).all()
+        return bool(
+            (other.betas_ == self.betas_).all()
+            and (other.X_train == self.X_train).all()
+        )

legateboost/models/linear.py

@@ -121,4 +121,4 @@ def __str__(self) -> str:
     def __eq__(self, other: object) -> bool:
         if not isinstance(other, Linear):
             raise NotImplementedError()
-        return (other.betas_ == self.betas_).all()
+        return bool((other.betas_ == self.betas_).all())

legateboost/objectives.py

@@ -21,6 +21,17 @@
 
 GradPair: TypeAlias = Tuple[cn.ndarray, cn.ndarray]
 
+__all__ = [
+    "BaseObjective",
+    "SquaredErrorObjective",
+    "NormalObjective",
+    "LogLossObjective",
+    "ExponentialObjective",
+    "QuantileObjective",
+    "GammaDevianceObjective",
+    "GammaObjective",
+]
+
 
 class BaseObjective(ABC):
     """The base class for objective functions.

legateboost/shapley.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Optional, Tuple
+from typing import TYPE_CHECKING, List, Optional, Tuple
 
 import numpy as np
 from sklearn.base import is_regressor
@@ -12,6 +12,8 @@
 if TYPE_CHECKING:
     from .legateboost import LBBase
 
+__all__: List[str] = []
+
 
 def global_shapley_attributions(
     model: "LBBase",

legateboost/special.py

@@ -15,13 +15,16 @@
 from __future__ import annotations
 
 from enum import IntEnum
+from typing import List
 
 import cunumeric as cn
 from legate.core import get_legate_runtime, types as ty
 
 from .library import user_context, user_lib
 from .utils import get_store
 
+__all__: List[str] = []
+
 
 class _SpecialOpCode(IntEnum):
     ERF = user_lib.cffi.ERF

legateboost/utils.py

@@ -14,6 +14,8 @@
 
 from .library import user_context, user_lib
 
+__all__: List[str] = []
+
 
 class PickleCunumericMixin:
     """When reading back from pickle, convert numpy arrays to cunumeric