Add some coverage

Author: CeliaBenquet

cebra/solver/base.py

@@ -360,6 +360,12 @@ def _get_loader(self, loader):
 
     @abc.abstractmethod
     def _set_fitted_params(self, loader: cebra.data.Loader):
+        """Set parameters once the solver is fitted.
+
+        Args:
+            loader: Loader used to fit the solver.
+        """
+
         raise NotImplementedError
 
     def fit(
@@ -507,6 +513,11 @@ def _check_is_inputs_valid(self, inputs: torch.Tensor, session_id: int):
 
     @abc.abstractmethod
     def _check_is_session_id_valid(self, session_id: Optional[int] = None):
+        """Check that the session ID provided is valid for the solver instance.
+        
+        Args: 
+            session_id: The session ID to check.
+        """
         raise NotImplementedError
 
     @abc.abstractmethod
@@ -530,7 +541,7 @@ def _select_model(
 
     @torch.no_grad()
     def transform(self,
-                  inputs: torch.Tensor,
+                  inputs: Union[torch.Tensor, List[torch.Tensor], npt.NDArray],
                   pad_before_transform: bool = True,
                   session_id: Optional[int] = None,
                   batch_size: Optional[int] = None) -> torch.Tensor:

cebra/solver/multi_session.py

@@ -126,6 +126,17 @@ def _inference(self, batches: List[cebra.data.Batch]) -> cebra.data.Batch:
         )
 
     def _set_fitted_params(self, loader: cebra.data.Loader):
+        """Set parameters once the solver is fitted.
+        
+        In multi session solver, the number of session is set to the number of
+        sessions in the dataset of the loader and the number of
+        features is set as a list corresponding to the number of neurons in 
+        each dataset.
+
+        Args:
+            loader: Loader used to fit the solver.
+        """
+
         self.num_sessions = loader.dataset.num_sessions
         self.n_features = [
             loader.dataset.get_input_dimension(session_id)
@@ -152,6 +163,14 @@ def _check_is_inputs_valid(self, inputs: torch.Tensor,
             )
 
     def _check_is_session_id_valid(self, session_id: Optional[int]):
+        """Check that the session ID provided is valid for the solver instance.
+        
+        The session ID must be non-null and between 0 and the number session in the dataset.
+        
+        Args: 
+            session_id: The session ID to check.
+        """
+
         if session_id is None:
             raise RuntimeError(
                 "No session_id provided: multisession model requires a session_id to choose the model corresponding to your data shape."

cebra/solver/single_session.py

@@ -55,6 +55,14 @@ def parameters(self, session_id: Optional[int] = None):
             yield parameter
 
     def _set_fitted_params(self, loader: cebra.data.Loader):
+        """Set parameters once the solver is fitted.
+        
+        In single session solver, the number of session is set to None and the number of
+        features is set to the number of neurons in the dataset.
+
+        Args:
+            loader: Loader used to fit the solver.
+        """
         self.num_sessions = None
         self.n_features = loader.dataset.input_dimension
 
@@ -77,6 +85,14 @@ def _check_is_inputs_valid(self, inputs: torch.Tensor, session_id: int):
             )
 
     def _check_is_session_id_valid(self, session_id: Optional[int] = None):
+        """Check that the session ID provided is valid for the solver instance.
+        
+        The session ID must be null or equal to 0.
+        
+        Args: 
+            session_id: The session ID to check.
+        """
+
         if session_id is not None and session_id > 0:
             raise RuntimeError(
                 f"Invalid session_id {session_id}: single session models only takes an optional null session_id."