Add prod to BatchedTensor

pyproject.toml

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

src/redcat/tensor.py

@@ -2184,6 +2184,37 @@ def nansum(self, *args, **kwargs) -> Tensor:
         """
         return torch.nansum(self, *args, **kwargs)
 
+    def prod(self, *args, **kwargs) -> Tensor:
+        r"""Computes the product of all elements.
+
+        Args:
+            *args: See the documentation of ``torch.Tensor.nansum``
+            **kwargs: See the documentation of ``torch.Tensor.nansum``
+
+        Returns:
+            ``torch.Tensor``: The product of all elements.
+
+        Example usage:
+
+        .. code-block:: python
+
+            >>> import torch
+            >>> from redcat import BatchedTensor
+            >>> BatchedTensor(
+            ...     torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 1]])
+            ... ).prod()
+            tensor(362880)
+            >>> BatchedTensor(
+            ...     torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 1]])
+            ... ).prod(dim=1)
+            tensor([ 120, 3024])
+            >>> BatchedTensor(
+            ...     torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 1]])
+            ... ).prod(dim=1, keepdim=True)
+            tensor([[ 120], [3024]])
+        """
+        return torch.prod(self, *args, **kwargs)
+
     def sum(self, *args, **kwargs) -> Tensor:
         r"""Computes the sum of all elements.
 

tests/unit/test_tensor.py

@@ -2735,6 +2735,38 @@ def test_batched_tensor_nansum_custom_dims() -> None:
     )
 
 
+def test_batched_tensor_prod() -> None:
+    assert (
+        BatchedTensor(torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 1]]))
+        .prod()
+        .equal(torch.tensor(362880))
+    )
+
+
+def test_batched_tensor_prod_keepdim_false() -> None:
+    assert (
+        BatchedTensor(torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 1]]))
+        .prod(dim=1)
+        .equal(torch.tensor([120, 3024]))
+    )
+
+
+def test_batched_tensor_prod_keepdim_true() -> None:
+    assert (
+        BatchedTensor(torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 1]]))
+        .prod(dim=1, keepdim=True)
+        .equal(torch.tensor([[120], [3024]]))
+    )
+
+
+def test_batched_tensor_prod_custom_dims() -> None:
+    assert (
+        BatchedTensor(torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 1]]), batch_dim=1)
+        .prod()
+        .equal(torch.tensor(362880))
+    )
+
+
 @mark.parametrize("dtype", (torch.float, torch.long))
 def test_batched_tensor_sum(dtype: torch.dtype) -> None:
     assert (

tests/unit/test_tensorseq.py

@@ -5023,6 +5023,38 @@ def test_batched_tensor_seq_nansum_custom_dims() -> None:
     )
 
 
+def test_batched_tensor_seq_prod() -> None:
+    assert (
+        BatchedTensorSeq(torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 1]]))
+        .prod()
+        .equal(torch.tensor(362880))
+    )
+
+
+def test_batched_tensor_seq_prod_keepdim_false() -> None:
+    assert (
+        BatchedTensorSeq(torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 1]]))
+        .prod(dim=1)
+        .equal(torch.tensor([120, 3024]))
+    )
+
+
+def test_batched_tensor_seq_prod_keepdim_true() -> None:
+    assert (
+        BatchedTensorSeq(torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 1]]))
+        .prod(dim=1, keepdim=True)
+        .equal(torch.tensor([[120], [3024]]))
+    )
+
+
+def test_batched_tensor_seq_prod_custom_dims() -> None:
+    assert (
+        BatchedTensorSeq(torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 1]]), batch_dim=1, seq_dim=0)
+        .prod()
+        .equal(torch.tensor(362880))
+    )
+
+
 @mark.parametrize("dtype", (torch.float, torch.long))
 def test_batched_tensor_seq_sum(dtype: torch.dtype) -> None:
     assert (