Merge branch 'main' into batched-inference-and-padding

stes · web-flow · commit b1980cdebcf2 · 2025-01-21T23:54:12.000+01:00
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
@@ -19,10 +19,16 @@ jobs:
         # as well as selected previous versions on
         # https://pytorch.org/get-started/previous-versions/
         torch-version: ["2.2.2", "2.4.0"]
+        sklearn-version: ["latest"]
         include:
           - os: windows-latest
             torch-version: 2.4.0
             python-version: "3.10"
+            sklearn-version: "latest"
+          - os: ubuntu-latest 
+            torch-version: 2.4.0
+            python-version: "3.10"
+            sklearn-version: "legacy"
 
     runs-on: ${{ matrix.os }}
 
@@ -32,7 +38,7 @@ jobs:
         uses: actions/cache@v3
         with:
           path: ~/.cache/pip
-          key: pip-os_${{ runner.os }}-python_${{ matrix.python-version }}-torch_${{ matrix.torch-version }}
+          key: pip-os_${{ runner.os }}-python_${{ matrix.python-version }}-torch_${{ matrix.torch-version }}-sklearn_${{ matrix.sklearn-version }}
 
       - name: Checkout code
         uses: actions/checkout@v2
@@ -48,6 +54,11 @@ jobs:
           python -m pip install torch==${{ matrix.torch-version }} --extra-index-url https://download.pytorch.org/whl/cpu
           pip install '.[dev,datasets,integrations]'
 
+      - name: Check sklearn legacy version 
+        if: matrix.sklearn-version == 'legacy'
+        run: |
+          pip install scikit-learn==1.4.2 '.[dev,datasets,integrations]'
+
       - name: Run the formatter
         run: |
           make format
diff --git a/cebra/integrations/sklearn/cebra.py b/cebra/integrations/sklearn/cebra.py
@@ -30,8 +30,10 @@
 import pkg_resources
 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
 from torch import nn
 
 import cebra.data
@@ -41,6 +43,11 @@
 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,
@@ -364,7 +371,7 @@ def _load_cebra_with_sklearn_backend(cebra_info: Dict) -> "CEBRA":
     return cebra_
 
 
-class CEBRA(BaseEstimator, TransformerMixin):
+class CEBRA(TransformerMixin, BaseEstimator):
     """CEBRA model defined as part of a ``scikit-learn``-like API.
 
     Attributes:
@@ -1317,6 +1324,15 @@ def fit_transform(
                  callback_frequency=callback_frequency)
         return self.transform(X)
 
+    def __sklearn_tags__(self):
+        # NOTE(stes): from 1.6.dev, this is the new way to specify tags
+        # https://scikit-learn.org/dev/developers/develop.html
+        # https://github.com/scikit-learn/scikit-learn/pull/29677#issuecomment-2334229165
+        tags = super().__sklearn_tags__()
+        tags.non_deterministic = True
+        return tags
+
+    @available_if(check_version)
     def _more_tags(self):
         # NOTE(stes): This tag is needed as seeding is not fully implemented in the
         # current version of CEBRA.
diff --git a/conda/cebra_paper.yml b/conda/cebra_paper.yml
@@ -39,7 +39,7 @@ dependencies:
         - "cebra[dev,integrations,datasets,demos]"
         - joblib
         - literate-dataclasses
-        - sklearn
+        - scikit-learn
         - scipy
         - torch
         - keras==2.3.1
diff --git a/conda/cebra_paper_m1.yml b/conda/cebra_paper_m1.yml
@@ -48,7 +48,7 @@ dependencies:
         - tensorflow-metal
         - joblib
         - literate-dataclasses
-        - sklearn
+        - scikit-learn
         - scipy
         - torch
         - umap-learn
diff --git a/tests/test_api.py b/tests/test_api.py
@@ -21,6 +21,5 @@
 #
 def test_api():
     import cebra.distributions
-    from cebra.distributions import TimedeltaDistribution
 
     cebra.distributions.TimedeltaDistribution
diff --git a/tests/test_cli.py b/tests/test_cli.py
@@ -19,6 +19,3 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-import argparse
-
-import pytest
diff --git a/tests/test_criterions.py b/tests/test_criterions.py
@@ -19,7 +19,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-import numpy as np
 import pytest
 import torch
 from torch import nn
@@ -294,7 +293,7 @@ def _sample_dist_matrices(seed):
 
 
 @pytest.mark.parametrize("seed", [42, 4242, 424242])
-def test_infonce(seed):
+def test_infonce_check_output_parts(seed):
     pos_dist, neg_dist = _sample_dist_matrices(seed)
 
     ref_loss, ref_align, ref_uniform = _reference_infonce(pos_dist, neg_dist)
diff --git a/tests/test_datasets.py b/tests/test_datasets.py
@@ -99,7 +99,6 @@ def test_hippocampus():
 
 @pytest.mark.requires_dataset
 def test_monkey():
-
     dataset = cebra.datasets.init(
         "area2-bump-pos-active-passive",
         path=pathlib.Path(_DEFAULT_DATADIR) / "monkey_reaching_preload_smth_40",
@@ -110,7 +109,6 @@ def test_monkey():
 
 @pytest.mark.requires_dataset
 def test_allen():
-
     pytest.skip("Test takes too long")
 
     ca_dataset = cebra.datasets.init("allen-movie-one-ca-VISp-100-train-10-111")
diff --git a/tests/test_demo.py b/tests/test_demo.py
@@ -21,7 +21,6 @@
 #
 import glob
 import re
-import sys
 
 import pytest
 
diff --git a/tests/test_distributions.py b/tests/test_distributions.py
@@ -43,7 +43,7 @@ def prepare(N=1000, n=128, d=5, probs=[0.3, 0.1, 0.6], device="cpu"):
     continuous = torch.randn(N, d).to(device)
 
     rand = torch.from_numpy(np.random.randint(0, N, (n,))).to(device)
-    qidx = discrete[rand].to(device)
+    _ = discrete[rand].to(device)
     query = continuous[rand] + 0.1 * torch.randn(n, d).to(device)
     query = query.to(device)
 
@@ -173,7 +173,7 @@ def test_mixed():
         discrete, continuous)
 
     reference_idx = distribution.sample_prior(10)
-    positive_idx = distribution.sample_conditional(reference_idx)
+    _ = distribution.sample_conditional(reference_idx)
 
     # The conditional distribution p(· | disc, cont) should yield
     # samples where the label exactly matches the reference sample.
@@ -193,7 +193,7 @@ def test_continuous(benchmark):
     def _test_distribution(dist):
         distribution = dist(continuous)
         reference_idx = distribution.sample_prior(10)
-        positive_idx = distribution.sample_conditional(reference_idx)
+        _ = distribution.sample_conditional(reference_idx)
         return distribution
 
     distribution = _test_distribution(
diff --git a/tests/test_grid_search.py b/tests/test_grid_search.py
@@ -20,7 +20,6 @@
 # limitations under the License.
 #
 import numpy as np
-import pytest
 
 import cebra
 import cebra.grid_search
diff --git a/tests/test_integration_train.py b/tests/test_integration_train.py
@@ -20,7 +20,6 @@
 # limitations under the License.
 #
 import itertools
-from typing import List
 
 import pytest
 import torch
diff --git a/tests/test_load.py b/tests/test_load.py
@@ -22,10 +22,7 @@
 import itertools
 import pathlib
 import pickle
-import platform
 import tempfile
-import unittest
-from unittest.mock import patch
 
 import h5py
 import hdf5storage
@@ -125,7 +122,7 @@ def generate_numpy_confounder(filename, dtype):
 
 
 @register("npz")
-def generate_numpy_path(filename, dtype):
+def generate_numpy_path_2(filename, dtype):
     A = np.arange(1000, dtype=dtype).reshape(10, 100)
     np.savez(filename, array=A, other_data="test")
     loaded_A = cebra_load.load(pathlib.Path(filename))
@@ -418,7 +415,7 @@ def generate_csv_path(filename, dtype):
 
 @register_error("csv")
 def generate_csv_empty_file(filename, dtype):
-    with open(filename, "w") as creating_new_csv_file:
+    with open(filename, "w") as _:
         pass
     _ = cebra_load.load(filename)
 
@@ -619,7 +616,6 @@ def generate_pickle_invalid_key(filename, dtype):
 
 @register_error("pkl", "p")
 def generate_pickle_no_array(filename, dtype):
-    A = np.arange(1000, dtype=dtype).reshape(10, 100)
     with open(filename, "wb") as f:
         pickle.dump({"A": "test_1", "B": "test_2"}, f)
     _ = cebra_load.load(filename)
diff --git a/tests/test_models.py b/tests/test_models.py
@@ -155,8 +155,8 @@ def test_version_check(version, raises):
             cebra.models.model._check_torch_version(raise_error=True)
 
 
-def test_version_check():
-    raises = not cebra.models.model._check_torch_version(raise_error=False)
+def test_version_check_dropout_available():
+    raises = cebra.models.model._check_torch_version(raise_error=False)
     if raises:
         assert len(cebra.models.get_options("*dropout*")) == 0
     else:
diff --git a/tests/test_plot.py b/tests/test_plot.py
@@ -72,8 +72,6 @@ def test_plot_imports():
 def test_colormaps():
     import matplotlib
 
-    import cebra
-
     cmap = matplotlib.colormaps["cebra"]
     assert cmap is not None
     plt.scatter([1], [2], c=[2], cmap="cebra")
@@ -241,7 +239,7 @@ def test_compare_models():
         _ = cebra_plot.compare_models(models, labels=long_labels, ax=ax)
     with pytest.raises(ValueError, match="Invalid.*labels"):
         invalid_labels = copy.deepcopy(labels)
-        ele = invalid_labels.pop()
+        _ = invalid_labels.pop()
         invalid_labels.append(["a"])
         _ = cebra_plot.compare_models(models, labels=invalid_labels, ax=ax)
 
diff --git a/tests/test_registry.py b/tests/test_registry.py
@@ -117,7 +117,7 @@ def test_override():
     _Foo1 = test_module.register("foo")(Foo)
     assert _Foo1 == Foo
     assert _Foo1 != Bar
-    assert f"foo" in test_module.get_options()
+    assert "foo" in test_module.get_options()
 
     # Check that the class was actually added to the module
     assert (
@@ -137,15 +137,15 @@ def test_override():
     _Foo2 = test_module.register("foo", override=True)(Bar)
     assert _Foo2 != Foo
     assert _Foo2 == Bar
-    assert f"foo" in test_module.get_options()
+    assert "foo" in test_module.get_options()
 
 
 def test_depreciation():
     test_module = _make_registry()
     Foo = _make_class()
     _Foo1 = test_module.register("foo")(Foo)
     assert _Foo1 == Foo
-    assert f"foo" in test_module.get_options()
+    assert "foo" in test_module.get_options()
 
     # Registering the same class under different names
     # also raises and error
diff --git a/tests/test_sklearn.py b/tests/test_sklearn.py
@@ -276,7 +276,6 @@ def test_api(estimator, check):
         pytest.skip(f"Model architecture {estimator.model_architecture} "
                     f"requires longer input sizes than 20 samples.")
 
-    success = True
     exception = None
     num_successful = 0
     total_runs = 0
@@ -334,7 +333,6 @@ def test_sklearn(model_architecture, device):
     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,))
 
@@ -863,7 +861,6 @@ def test_sklearn_full(model_architecture, device, pad_before_transform):
     X = np.random.uniform(0, 1, (1000, 50))
     y_c1 = np.random.uniform(0, 1, (1000, 5))
     y_c2 = np.random.uniform(0, 1, (1000, 2))
-    y_d = np.random.randint(0, 10, (1000,))
 
     # time contrastive
     cebra_model.fit(X)
@@ -931,7 +928,7 @@ def test_sklearn_resampling_model_not_yet_supported(model_architecture, device):
 
     with pytest.raises(ValueError):
         cebra_model.fit(X, y_c1)
-        output = cebra_model.transform(X)
+        _ = cebra_model.transform(X)
 
 
 def _iterate_actions():
diff --git a/tests/test_solver.py b/tests/test_solver.py
@@ -199,11 +199,12 @@ def test_single_session(data_name, loader_initfunc, model_architecture,
     single_session_tests)
 def test_single_session_auxvar(data_name, loader_initfunc, model_architecture,
                                solver_initfunc):
-    return  # TODO
+
+    pytest.skip("Not yet supported")
 
     loader = _get_loader(data_name, loader_initfunc)
     model = _make_model(loader.dataset)
-    behavior_model = _make_behavior_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)
diff --git a/tests/test_usecases.py b/tests/test_usecases.py
@@ -29,7 +29,6 @@
 """
 
 import itertools
-import pickle
 
 import numpy as np
 import pytest

Original file line number	Diff line number	Diff line change
`@@ -21,6 +21,5 @@`
`21`	`21`	`#`
`22`	`22`	`def test_api():`
`23`	`23`	`import cebra.distributions`
`24`		`- from cebra.distributions import TimedeltaDistribution`
`25`	`24`
`26`	`25`	`cebra.distributions.TimedeltaDistribution`
Original file line number	Diff line number	Diff line change
`@@ -19,6 +19,3 @@`
`19`	`19`	`# See the License for the specific language governing permissions and`
`20`	`20`	`# limitations under the License.`
`21`	`21`	`#`
`22`		`-import argparse`
`23`		`-`
`24`		`-import pytest`
Original file line number	Diff line number	Diff line change
`@@ -21,7 +21,6 @@`
`21`	`21`	`#`
`22`	`22`	`import glob`
`23`	`23`	`import re`
`24`		`-import sys`
`25`	`24`
`26`	`25`	`import pytest`
`27`	`26`