Update tests and duplicate code based on review

Author: CeliaBenquet

cebra/integrations/sklearn/cebra.py

@@ -830,6 +830,8 @@ def _configure_for_all(
 
     def _select_model(self, X: Union[npt.NDArray, torch.Tensor],
                       session_id: int):
+        if isinstance(X, np.ndarray):
+            X = torch.from_numpy(X)
         return self.solver_._select_model(X, session_id=session_id)
 
     def _check_labels_types(self, y: tuple, session_id: Optional[int] = None):
@@ -1061,7 +1063,8 @@ def _partial_fit(
         self.model_ = model
 
         self.n_features_ = solver.n_features
-        self.num_sessions_ = solver.num_sessions
+        self.num_sessions_ = solver.num_sessions if hasattr(
+            solver, "num_sessions") else None
         self.solver_ = solver
         self.n_features_in_ = ([model[n].num_input for n in range(len(model))]
                                if is_multisession else model.num_input)
@@ -1247,10 +1250,6 @@ def transform(self,
         if isinstance(X, np.ndarray):
             X = torch.from_numpy(X)
 
-        if batch_size is not None and batch_size < 1:
-            raise ValueError(
-                f"Batch size should be at least 1, got {batch_size}")
-
         with torch.no_grad():
             output = self.solver_.transform(
                 inputs=X,

cebra/solver/base.py

@@ -36,7 +36,6 @@
 from typing import Callable, Dict, List, Literal, Optional, Tuple, Union
 
 import literate_dataclasses as dataclasses
-import numpy.typing as npt
 import torch
 import torch.nn.functional as F
 from torch.utils.data import DataLoader
@@ -104,6 +103,15 @@ def _add_batched_zero_padding(batched_data: torch.Tensor,
     Returns:
         The padded batch.
     """
+    if batch_start_idx > batch_end_idx:
+        raise ValueError(
+            f"batch_start_idx ({batch_start_idx}) cannot be greater than batch_end_idx ({batch_end_idx})."
+        )
+    if batch_start_idx < 0 or batch_end_idx < 0:
+        raise ValueError(
+            f"batch_start_idx ({batch_start_idx}) and batch_end_idx ({batch_end_idx}) must be positive integers."
+        )
+
     reversed_dims = torch.arange(batched_data.ndim - 1, -1, -1)
 
     if batch_start_idx == 0:  # First batch
@@ -179,7 +187,7 @@ def _inference_transform(model: cebra.models.Model,
     return output
 
 
-def _transform(
+def _not_batched_transform(
     model: cebra.models.Model,
     inputs: torch.Tensor,
     pad_before_transform: bool,
@@ -253,9 +261,11 @@ def __getitem__(self, idx):
         if batch_idx == (len(index_dataloader) - 1):
             # last batch, incomplete
             index_batch = torch.cat((last_batch, index_batch), dim=0)
-            assert index_batch[-1] + 1 == len(inputs), (
-                f"Last batch index {index_batch[-1]} + 1 should be equal to the length of inputs {len(inputs)}."
-            )
+
+            if index_batch[-1] + 1 != len(inputs):
+                raise ValueError(
+                    f"Last batch index {index_batch[-1]} + 1 should be equal to the length of inputs {len(inputs)}."
+                )
 
         # Batch start and end so that `batch_size` size with the last batch including 2 batches
         batch_start_idx, batch_end_idx = index_batch[0], index_batch[-1] + 1
@@ -494,9 +504,6 @@ def fit(
                 if logdir is not None:
                     self.save(logdir, f"checkpoint_{num_steps:#07d}.pth")
 
-        assert hasattr(self, "n_features")
-        assert hasattr(self, "num_sessions")
-
     def step(self, batch: cebra.data.Batch) -> dict:
         """Perform a single gradient update.
 
@@ -540,7 +547,10 @@ def validation(self,
         Returns:
             Loss averaged over iterations on data batch.
         """
-        assert (session_id is None) or (session_id == 0)
+        if session_id is not None and session_id != 0:
+            raise ValueError(
+                f"session_id should be set to None or 0, got {session_id}")
+
         iterator = self._get_loader(loader)
         total_loss = Meter()
         self.model.eval()
@@ -569,7 +579,6 @@ def decoding(self, train_loader, valid_loader):
         )
         return decode_metric
 
-    @abc.abstractmethod
     def _check_is_inputs_valid(self, inputs: torch.Tensor, session_id: int):
         """Check that the inputs can be inferred using the selected model.
 
@@ -582,7 +591,13 @@ def _check_is_inputs_valid(self, inputs: torch.Tensor, session_id: int):
                 the number of sessions -1 for multisession, and set to
                 ``None`` for single session.
         """
-        raise NotImplementedError
+        if isinstance(inputs, list):
+            raise ValueError(
+                "Inputs to transform() should be the data for a single session, but received a list."
+            )
+        elif not isinstance(inputs, torch.Tensor):
+            raise ValueError(
+                f"Inputs should be a torch.Tensor, not {type(inputs)}.")
 
     @abc.abstractmethod
     def _check_is_session_id_valid(self, session_id: Optional[int] = None):
@@ -593,7 +608,6 @@ def _check_is_session_id_valid(self, session_id: Optional[int] = None):
         """
         raise NotImplementedError
 
-    @abc.abstractmethod
     def _select_model(
         self, inputs: Union[torch.Tensor,
                             List[torch.Tensor]], session_id: Optional[int]
@@ -610,6 +624,25 @@ def _select_model(
         Returns:
             The model (first returns) and the offset of the model (second returns).
         """
+        model = self._get_model(session_id=session_id)
+        offset = model.get_offset()
+
+        self._check_is_inputs_valid(inputs, session_id=session_id)
+        return model, offset
+
+    @abc.abstractmethod
+    def _get_model(self,
+                   session_id: Optional[int] = None) -> cebra.models.Model:
+        """Get the model to use for inference.
+
+        Args:
+            session_id: The session ID, an :py:class:`int` between 0 and
+                the number of sessions -1 for multisession, and set to
+                ``None`` for single session.
+
+        Returns:
+            The model.
+        """
         raise NotImplementedError
 
     def _check_is_fitted(self):
@@ -627,7 +660,7 @@ def _check_is_fitted(self):
 
     @torch.no_grad()
     def transform(self,
-                  inputs: Union[torch.Tensor, List[torch.Tensor], npt.NDArray],
+                  inputs: torch.Tensor,
                   pad_before_transform: Optional[bool] = True,
                   session_id: Optional[int] = None,
                   batch_size: Optional[int] = None) -> torch.Tensor:
@@ -652,26 +685,40 @@ def transform(self,
         Returns:
             The output embedding.
         """
-        if isinstance(inputs, list):
-            raise ValueError(
-                "Inputs to transform() should be the data for a single session, but received a list."
-            )
-        elif not isinstance(inputs, torch.Tensor):
-            raise ValueError(
-                f"Inputs should be a torch.Tensor, not {type(inputs)}.")
-
         self._check_is_fitted()
-
         model, offset = self._select_model(inputs, session_id)
 
         if len(offset) < 2 and pad_before_transform:
             pad_before_transform = False
 
         model.eval()
+        return self._transform(model=model,
+                               inputs=inputs,
+                               pad_before_transform=pad_before_transform,
+                               offset=offset,
+                               batch_size=batch_size)
+
+    @torch.no_grad()
+    def _transform(self, model: cebra.models.Model, inputs: torch.Tensor,
+                   pad_before_transform: bool,
+                   offset: cebra.data.datatypes.Offset,
+                   batch_size: Optional[int]) -> torch.Tensor:
+        """Compute the embedding on the inputs using the model provided.
+
+        Args:
+            model: Model to use for inference.
+            inputs: Data.
+            pad_before_transform: If True zero-pad the batched data.
+            offset: Offset of the model to consider when padding.
+            batch_size: If not None, batched inference will not be applied.
+
+        Returns:
+            The embedding.
+        """
         if batch_size is not None and inputs.shape[0] > int(
-                batch_size * 2) and not isinstance(
-                    self.model, cebra.models.ResampleModelMixin):
-            # NOTE: resampling models are not supported for batched inference.
+                batch_size * 2) and not (isinstance(
+                    self._get_model(0), cebra.models.ResampleModelMixin)):
+            # NOTE(celia): resampling models are not supported for batched inference.
             output = _batched_transform(
                 model=model,
                 inputs=inputs,
@@ -680,11 +727,11 @@ def transform(self,
                 pad_before_transform=pad_before_transform,
             )
         else:
-            output = _transform(model=model,
-                                inputs=inputs,
-                                offset=offset,
-                                pad_before_transform=pad_before_transform)
-
+            output = _not_batched_transform(
+                model=model,
+                inputs=inputs,
+                offset=offset,
+                pad_before_transform=pad_before_transform)
         return output
 
     @abc.abstractmethod
@@ -863,3 +910,37 @@ def step(self, batch: cebra.data.Batch) -> dict:
             time_neg=time_uniform.item(),
             time_total=time_loss.item(),
         )
+
+
+class AuxiliaryVariableSolver(Solver):
+
+    @torch.no_grad()
+    def transform(self,
+                  inputs: torch.Tensor,
+                  pad_before_transform: bool = True,
+                  session_id: Optional[int] = None,
+                  batch_size: Optional[int] = None,
+                  use_reference_model: bool = False) -> torch.Tensor:
+        """Compute the embedding.
+        This function by default use ``model`` that was trained to encode the positive
+        and negative samples. To use ``reference_model`` instead of ``model``
+        ``use_reference_model`` should be equal ``True``.
+        Args:
+            inputs: The input signal
+            use_reference_model: Flag for using ``reference_model``
+        Returns:
+            The output embedding.
+        """
+        self._check_is_fitted()
+        model, offset = self._select_model(
+            inputs, session_id, use_reference_model=use_reference_model)
+
+        if len(offset) < 2 and pad_before_transform:
+            pad_before_transform = False
+
+        model.eval()
+        return self._transform(model=model,
+                               inputs=inputs,
+                               pad_before_transform=pad_before_transform,
+                               offset=offset,
+                               batch_size=batch_size)