Start to implement BatchList

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "redcat"
-version = "0.0.1a194"
+version = "0.0.1a195"
 description = "A library to manipulate batches of examples"
 readme = "README.md"
 authors = ["Thibaut Durand <durand.tibo+gh@gmail.com>"]

src/redcat/__init__.py

@@ -1,6 +1,12 @@
-__all__ = ["BaseBatch", "BatchedTensor", "BatchedTensorSeq"]
+__all__ = [
+    "BaseBatch",
+    "BatchList",
+    "BatchedTensor",
+    "BatchedTensorSeq",
+]
 
 from redcat import comparators  # noqa: F401
 from redcat.base import BaseBatch
+from redcat.list import BatchList
 from redcat.tensor import BatchedTensor
 from redcat.tensorseq import BatchedTensorSeq

src/redcat/base.py

@@ -33,15 +33,11 @@ def data(self) -> T:
     ###############################
 
     @abstractmethod
-    def clone(self, *args, **kwargs) -> TBatch:
+    def clone(self) -> TBatch:
         r"""Creates a copy of the current batch.
 
-        Args:
-            *args: See the documentation of ``torch.Tensor.clone``
-            **kwargs: See the documentation of ``torch.Tensor.clone``
-
         Returns:
-            ``BaseBatchedTensor``: A copy of the current batch.
+            ``BaseBatch``: A copy of the current batch.
 
         Example usage:
 
@@ -56,6 +52,10 @@ def clone(self, *args, **kwargs) -> TBatch:
                     [1., 1., 1.]], batch_dim=0)
         """
 
+    #################################
+    #     Comparison operations     #
+    #################################
+
     @abstractmethod
     def allclose(
         self, other: Any, rtol: float = 1e-5, atol: float = 1e-8, equal_nan: bool = False
@@ -123,7 +123,7 @@ def permute_along_batch(self, permutation: Sequence[int] | Tensor) -> TBatch:
                 input should be compatible with the shape of the data.
 
         Returns:
-            ``BaseBatchedTensor``: A new batch with permuted data.
+            ``BaseBatch``: A new batch with permuted data.
 
         Example usage:
 
@@ -175,7 +175,7 @@ def shuffle_along_batch(self, generator: torch.Generator | None = None) -> TBatc
                 Default: ``None``
 
         Returns:
-            ``BaseBatchedTensor``:  A new batch with shuffled data.
+            ``BaseBatch``:  A new batch with shuffled data.
 
         Example usage: