Extend type checking to all float datatypes (#166)

* Make the type checking less sensitive

* Use existing type checking functions

* Add better typing

* Update cebra/data/datasets.py

* Update cebra/data/datasets.py

---------

Co-authored-by: Steffen Schneider <[email protected]>
Co-authored-by: Mackenzie Mathis <[email protected]>

Author: 3 people

cebra/data/datasets.py

@@ -24,7 +24,7 @@
 import abc
 import collections
 import types
-from typing import List, Tuple, Union
+from typing import List, Literal, Optional, Tuple, Union
 
 import literate_dataclasses as dataclasses
 import numpy as np
@@ -74,23 +74,40 @@ def __init__(self,
                  offset: int = 1,
                  device: str = "cpu"):
         super().__init__(device=device)
-        self.neural = self._to_tensor(neural, torch.FloatTensor).float()
-        self.continuous = self._to_tensor(continuous, torch.FloatTensor)
-        self.discrete = self._to_tensor(discrete, torch.LongTensor)
+        self.neural = self._to_tensor(neural, check_dtype="float").float()
+        self.continuous = self._to_tensor(continuous,
+                                          check_dtype="float")
+        self.discrete = self._to_tensor(discrete, check_dtype="integer")
         if self.continuous is None and self.discrete is None:
             raise ValueError(
                 "You have to pass at least one of the arguments 'continuous' or 'discrete'."
             )
         self.offset = offset
 
-    def _to_tensor(self, array, check_dtype=None):
+    def _to_tensor(
+            self,
+            array: Union[torch.Tensor, npt.NDArray],
+            check_dtype: Optional[Literal["int",
+                                          "float"]] = None) -> torch.Tensor:
+        """Convert :py:func:`numpy.array` to :py:class:`torch.Tensor` if necessary and check the dtype.
+
+        Args:
+            array: Array to check.
+            check_dtype (list, optional): If not `None`, list of dtypes to which the values in `array`
+                must belong to. Defaults to None.
+
+        Returns:
+            The `array` as a :py:class:`torch.Tensor`.
+        """
         if array is None:
             return None
         if isinstance(array, np.ndarray):
             array = torch.from_numpy(array)
         if check_dtype is not None:
-            if not isinstance(array, check_dtype):
-                raise TypeError(f"{type(array)} instead of {check_dtype}.")
+            if (check_dtype == "int" and not cebra.helper._is_integer(array)
+               ) or (check_dtype == "float" and
+                     not cebra.helper._is_floating(array)):
+                raise TypeError(f"Array has type {array.dtype} instead of {check_dtype}.")
         return array
 
     @property

cebra/helper.py

@@ -99,7 +99,7 @@ def _is_integer(y: Union[npt.NDArray, torch.Tensor]) -> bool:
 
 
 def _is_floating(y: Union[npt.NDArray, torch.Tensor]) -> bool:
-    """Check if the values in ``y`` are :py:class:`int`.
+    """Check if the values in ``y`` are :py:class:`float`.
 
     Note:
         There is no ``torch`` method to check that the ``dtype`` of a :py:class:`torch.Tensor`
@@ -118,6 +118,19 @@ def _is_floating(y: Union[npt.NDArray, torch.Tensor]) -> bool:
                 y, torch.Tensor) and torch.is_floating_point(y))
 
 
+def _is_floating_or_integer(y: Union[npt.NDArray, torch.Tensor]) -> bool:
+    """Check if the values in ``y`` are :py:class:`int` or :py:class:`float`.
+
+    Args:
+        y: An array, either as a :py:func:`numpy.array` or a :py:class:`torch.Tensor`.
+
+    Returns:
+        ``True`` if ``y`` contains :py:class:`float` or :py:class:`int`.
+    """
+
+    return _is_floating(y) or _is_integer(y)
+
+
 def get_loader_options(dataset: "cebra.data.Dataset") -> List[str]:
     """Return all possible dataloaders for the given dataset.