Add align_to_seq_batch to BatchedTensorSeq

pyproject.toml

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

src/redcat/tensor_seq.py

@@ -13,6 +13,7 @@
 from redcat.tensor import BatchedTensor
 from redcat.utils import (
     align_to_batch_seq,
+    align_to_seq_batch,
     check_batch_dims,
     check_seq_dims,
     get_batch_dims,
@@ -697,9 +698,9 @@ def align_to_batch_seq(self) -> BatchedTensorSeq:
             >>> import torch
             >>> from redcat import BatchedTensorSeq
             >>> batch = BatchedTensorSeq(torch.arange(10).view(5, 2), batch_dim=1, seq_dim=0)
-            >>> batch.align_to_batch_seq().data
+            >>> batch.align_to_batch_seq()
             tensor([[0, 2, 4, 6, 8],
-                    [1, 3, 5, 7, 9]], batch_dim=0)
+                    [1, 3, 5, 7, 9]], batch_dim=0, seq_dim=1)
         """
         return self.__class__(
             data=align_to_batch_seq(
@@ -709,6 +710,34 @@ def align_to_batch_seq(self) -> BatchedTensorSeq:
             seq_dim=1,
         )
 
+    def align_to_seq_batch(self) -> BatchedTensorSeq:
+        r"""Aligns the current batch to the sequence-batch format.
+
+        Returns:
+            ``BatchedTensorSeq``: The batch in the sequence-batch format.
+
+        Example usage:
+
+        .. code-block:: python
+
+            >>> import torch
+            >>> from redcat import BatchedTensorSeq
+            >>> batch = BatchedTensorSeq(torch.arange(10).view(2, 5), batch_dim=0, seq_dim=1)
+            >>> batch.align_to_seq_batch()
+            tensor([[0, 5],
+                    [1, 6],
+                    [2, 7],
+                    [3, 8],
+                    [4, 9]], batch_dim=1, seq_dim=0)
+        """
+        return self.__class__(
+            data=align_to_seq_batch(
+                tensor=self._data, batch_dim=self._batch_dim, seq_dim=self._seq_dim
+            ),
+            batch_dim=1,
+            seq_dim=0,
+        )
+
     def cat_along_seq(
         self, other: BaseBatchedTensor | Tensor | Iterable[BaseBatchedTensor | Tensor]
     ) -> BatchedTensorSeq:

tests/unit/test_tensor_seq.py

@@ -4343,6 +4343,49 @@ def test_batched_tensor_seq_align_to_batch_seq_permute_dims_extra_dims() -> None
     )
 
 
+def test_batched_tensor_seq_align_to_seq_batch_no_permutation() -> None:
+    assert (
+        BatchedTensorSeq(torch.arange(10).view(2, 5), batch_dim=1, seq_dim=0)
+        .align_to_seq_batch()
+        .equal(
+            BatchedTensorSeq(
+                torch.tensor([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]), batch_dim=1, seq_dim=0
+            )
+        )
+    )
+
+
+def test_batched_tensor_seq_align_to_seq_batch_permute_dims() -> None:
+    assert (
+        BatchedTensorSeq(torch.arange(10).view(2, 5), batch_dim=0, seq_dim=1)
+        .align_to_seq_batch()
+        .equal(
+            BatchedTensorSeq(
+                torch.tensor([[0, 5], [1, 6], [2, 7], [3, 8], [4, 9]]), batch_dim=1, seq_dim=0
+            )
+        )
+    )
+
+
+def test_batched_tensor_seq_align_to_seq_batch_permute_dims_extra_dims() -> None:
+    assert (
+        BatchedTensorSeq(torch.arange(20).view(2, 5, 2), batch_dim=1, seq_dim=2)
+        .align_to_seq_batch()
+        .equal(
+            BatchedTensorSeq(
+                torch.tensor(
+                    [
+                        [[0, 10], [2, 12], [4, 14], [6, 16], [8, 18]],
+                        [[1, 11], [3, 13], [5, 15], [7, 17], [9, 19]],
+                    ]
+                ),
+                batch_dim=1,
+                seq_dim=0,
+            )
+        )
+    )
+
+
 @mark.parametrize(
     "other",
     (