Fix mypy issue

scripts/tests/test_build_model.py

@@ -21,7 +21,7 @@
 LIB_PATH = os.path.join(os.path.dirname(__file__), '..', '..')
 sys.path.insert(0, os.path.abspath(LIB_PATH))
 
-from scripts import build_model  # type: ignore # noqa (module hack)
+from scripts import build_model  # noqa (module hack)
 
 
 class TestAggregateScores(unittest.TestCase):

scripts/tests/test_encode_data.py

@@ -24,7 +24,7 @@
 LIB_PATH = os.path.join(os.path.dirname(__file__), '..', '..')
 sys.path.insert(0, os.path.abspath(LIB_PATH))
 
-from scripts import encode_data  # type: ignore # noqa (module hack)
+from scripts import encode_data  # noqa (module hack)
 
 
 class TestGetFeature(unittest.TestCase):

scripts/tests/test_prepare_knbc.py

@@ -21,7 +21,7 @@
 LIB_PATH = os.path.join(os.path.dirname(__file__), '..', '..')
 sys.path.insert(0, os.path.abspath(LIB_PATH))
 
-from scripts import prepare_knbc  # type: ignore # noqa (module hack)
+from scripts import prepare_knbc  # noqa (module hack)
 
 
 class TestBreakBeforeSequence(unittest.TestCase):

scripts/tests/test_train.py

@@ -27,7 +27,7 @@
 LIB_PATH = os.path.join(os.path.dirname(__file__), '..', '..')
 sys.path.insert(0, os.path.abspath(LIB_PATH))
 
-from scripts import train  # type: ignore # noqa (module hack)
+from scripts import train  # noqa (module hack)
 
 
 class TestArgParse(unittest.TestCase):

scripts/tests/test_translate_model.py

@@ -21,7 +21,7 @@
 LIB_PATH = os.path.join(os.path.dirname(__file__), '..', '..')
 sys.path.insert(0, os.path.abspath(LIB_PATH))
 
-from scripts import translate_model  # type: ignore # noqa (module hack)
+from scripts import translate_model  # noqa (module hack)
 
 
 class TestNormalize(unittest.TestCase):

scripts/train.py

@@ -157,10 +157,10 @@ def get_metrics(pred: npt.NDArray[np.bool_],
   Returns:
     result (Result): A result.
   """
-  tp = jnp.sum(jnp.logical_and(pred == 1, actual == 1))
-  tn = jnp.sum(jnp.logical_and(pred == 0, actual == 0))
-  fp = jnp.sum(jnp.logical_and(pred == 1, actual == 0))
-  fn = jnp.sum(jnp.logical_and(pred == 0, actual == 1))
+  tp: int = jnp.sum(jnp.logical_and(pred == 1, actual == 1))  # type: ignore
+  tn: int = jnp.sum(jnp.logical_and(pred == 0, actual == 0))  # type: ignore
+  fp: int = jnp.sum(jnp.logical_and(pred == 1, actual == 0))  # type: ignore
+  fn: int = jnp.sum(jnp.logical_and(pred == 0, actual == 1))  # type: ignore
   accuracy = (tp + tn) / (tp + tn + fp + fn)
   precision = tp / (tp + fp)
   recall = tp / (tp + fn)
@@ -208,7 +208,7 @@ def update(
   best_feature_index: int = err.argmin()
   positivity: bool = res.at[best_feature_index].get() < 0.5
   err_min = err.at[best_feature_index].get()
-  amount: float = jnp.log((1 - err_min) / (err_min + EPS))
+  amount: float = jnp.log((1 - err_min) / (err_min + EPS))  # type: ignore
 
   # This is equivalent to X_best = X[:, best_feature_index]
   X_best = jnp.zeros(

setup.cfg

@@ -28,7 +28,7 @@ dev =
     flake8
     isort
     numpy
-    mypy==0.971
+    mypy
     pytest
     toml
     twine