fix linting errors

Author: stes

cebra/integrations/sklearn/cebra.py

@@ -28,9 +28,9 @@
 import numpy as np
 import numpy.typing as npt
 import pkg_resources
+import sklearn
 import sklearn.utils.validation as sklearn_utils_validation
 import torch
-import sklearn
 from sklearn.base import BaseEstimator
 from sklearn.base import TransformerMixin
 from sklearn.utils.metaestimators import available_if
@@ -43,12 +43,14 @@
 import cebra.models
 import cebra.solver
 
+
 def check_version(estimator):
     # NOTE(stes): required as a check for the old way of specifying tags
     # https://github.com/scikit-learn/scikit-learn/pull/29677#issuecomment-2334229165
     from packaging import version
     return version.parse(sklearn.__version__) < version.parse("1.6.dev")
 
+
 def _init_loader(
     is_cont: bool,
     is_disc: bool,

tests/test_sklearn.py

@@ -1375,11 +1375,9 @@ def test_new_transform(model_architecture, device):
     # example dataset
     X = np.random.uniform(0, 1, (1000, 50))
     X_s2 = np.random.uniform(0, 1, (800, 30))
-    X_s3 = np.random.uniform(0, 1, (1000, 30))
     y_c1 = np.random.uniform(0, 1, (1000, 5))
     y_c1_s2 = np.random.uniform(0, 1, (800, 5))
     y_c2 = np.random.uniform(0, 1, (1000, 2))
-    y_c2_s2 = np.random.uniform(0, 1, (800, 2))
     y_d = np.random.randint(0, 10, (1000,))
     y_d_s2 = np.random.randint(0, 10, (800,))
 

tests/test_solver.py

@@ -163,7 +163,7 @@ def test_single_session(data_name, loader_initfunc, model_architecture,
 
     solver.fit(loader)
 
-    assert solver.num_sessions == None
+    assert solver.num_sessions is None
     assert solver.n_features == X.shape[1]
 
     embedding = solver.transform(X)
@@ -202,25 +202,25 @@ def test_single_session_auxvar(data_name, loader_initfunc, model_architecture,
 
     pytest.skip("Not yet supported")
 
-    loader = _get_loader(data_name, loader_initfunc)
-    model = _make_model(loader.dataset)
-    behavior_model = _make_behavior_model(loader.dataset)  # noqa: F841
+    # loader = _get_loader(data_name, loader_initfunc)
+    # model = _make_model(loader.dataset)
+    # behavior_model = _make_behavior_model(loader.dataset)  # noqa: F841
 
-    criterion = cebra.models.InfoNCE()
-    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
+    # criterion = cebra.models.InfoNCE()
+    # optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
 
-    solver = solver_initfunc(
-        model=model,
-        criterion=criterion,
-        optimizer=optimizer,
-    )
+    # solver = solver_initfunc(
+    #     model=model,
+    #     criterion=criterion,
+    #     optimizer=optimizer,
+    # )
 
-    batch = next(iter(loader))
-    assert batch.reference.shape == (32, loader.dataset.input_dimension, 10)
-    log = solver.step(batch)
-    assert isinstance(log, dict)
+    # batch = next(iter(loader))
+    # assert batch.reference.shape == (32, loader.dataset.input_dimension, 10)
+    # log = solver.step(batch)
+    # assert isinstance(log, dict)
 
-    solver.fit(loader)
+    # solver.fit(loader)
 
 
 @pytest.mark.parametrize(
@@ -251,7 +251,7 @@ def test_single_session_hybrid(data_name, loader_initfunc, model_architecture,
 
     solver.fit(loader)
 
-    assert solver.num_sessions == None
+    assert solver.num_sessions is None
     assert solver.n_features == X.shape[1]
 
     embedding = solver.transform(X)
@@ -513,7 +513,6 @@ def test_select_model_single_session(data_name, model_name, session_id,
     dataset = cebra.datasets.init(data_name)
     model = create_model(model_name, dataset.input_dimension)
     dataset.configure_for(model)
-    loader = _get_loader(dataset, loader_initfunc=loader_initfunc)
     offset = model.get_offset()
     solver = solver_initfunc(model=model, criterion=None, optimizer=None)
 
@@ -633,7 +632,6 @@ def test_batched_transform_single_session(
 
     smallest_batch_length = loader.dataset.neural.shape[0] - batch_size
     offset_ = model.get_offset()
-    padding_left = offset_.left if padding else 0
 
     if smallest_batch_length <= len(offset_):
         with pytest.raises(ValueError):
@@ -685,7 +683,6 @@ def test_batched_transform_multi_session(data_name, model_name, padding,
 
     smallest_batch_length = n_samples - batch_size
     offset_ = model[0].get_offset()
-    padding_left = offset_.left if padding else 0
     for d in dataset._datasets:
         d.offset = offset_
     loader_kwargs = dict(num_steps=10, batch_size=32)
@@ -710,7 +707,6 @@ def test_batched_transform_multi_session(data_name, model_name, padding,
                                  pad_before_transform=padding)
 
         else:
-            model_ = model[i]
             embedding = solver.transform(inputs=inputs.neural,
                                          session_id=i,
                                          pad_before_transform=padding)