Source formatting

Author: stes

cebra/integrations/sklearn/utils.py

@@ -22,22 +22,25 @@
 import warnings
 
 import numpy.typing as npt
+import packaging
+import sklearn
 import sklearn.utils.validation as sklearn_utils_validation
 import torch
 
 import cebra.helper
 
-from packaging import version
-import sklearn
 
 def _check_array_ensure_all_finite(array, **kwargs):
-    if version.parse(sklearn.__version__) < version.parse("1.8"):
+    # NOTE(stes): See discussion in https://github.com/AdaptiveMotorControlLab/CEBRA/pull/206
+    if packaging.version.parse(
+            sklearn.__version__) < packaging.version.parse("1.8"):
         key = "force_all_finite"
     else:
         key = "ensure_all_finite"
     kwargs[key] = True
     return sklearn_utils_validation.check_array(array, **kwargs)
 
+
 def update_old_param(old: dict, new: dict, kwargs: dict, default) -> tuple:
     """Handle deprecated arguments of a function until they are replaced.
 
@@ -85,17 +88,17 @@ def check_input_array(X: npt.NDArray, *, min_samples: int) -> npt.NDArray:
         The converted and validated array.
     """
     return _check_array_ensure_all_finite(
-          X,
-          accept_sparse=False,
-          accept_large_sparse=False,
-          dtype=("float16", "float32", "float64"),
-          order=None,
-          copy=False,
-          ensure_2d=True,
-          allow_nd=False,
-          ensure_min_samples=min_samples,
-          ensure_min_features=1,
-      )
+        X,
+        accept_sparse=False,
+        accept_large_sparse=False,
+        dtype=("float16", "float32", "float64"),
+        order=None,
+        copy=False,
+        ensure_2d=True,
+        allow_nd=False,
+        ensure_min_samples=min_samples,
+        ensure_min_features=1,
+    )
 
 
 def check_label_array(y: npt.NDArray, *, min_samples: int):