Add support for torch.nansum

pyproject.toml

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

src/redcat/tensor.py

@@ -3502,6 +3502,12 @@ def median(input: BatchedTensor, **kwargs) -> Tensor | torch.return_types.median
     return torch.median(input.data, **kwargs)
 
 
+@implements(torch.nansum)
+def nansum(input: BatchedTensor, **kwargs) -> Tensor:  # noqa: A002
+    r"""See ``torch.nansum`` documentation."""
+    return torch.nansum(input.data, **kwargs)
+
+
 @implements(torch.prod)
 def prod(input: BatchedTensor, **kwargs) -> Tensor:  # noqa: A002
     r"""See ``torch.prod`` documentation."""

src/redcat/tensorseq.py

@@ -24,6 +24,7 @@
 HANDLED_FUNCTIONS = {
     torch.mean: tensor.mean,
     torch.median: tensor.median,
+    torch.nansum: tensor.nansum,
     torch.prod: tensor.prod,
     torch.sum: tensor.torchsum,
 }

tests/unit/test_tensor.py

@@ -5324,6 +5324,31 @@ def test_torch_median_keepdim() -> None:
     )
 
 
+##################################
+#     Tests for torch.nansum     #
+##################################
+
+
+def test_torch_nansum() -> None:
+    assert torch.nansum(
+        BatchedTensor(torch.tensor([[0, 1, 2, 3, 4], [5, 6, 7, 8, float("nan")]]))
+    ).equal(torch.tensor(36))
+
+
+def test_torch_nansum_dim_1() -> None:
+    assert torch.nansum(
+        BatchedTensor(torch.tensor([[0, 1, 2, 3, 4], [5, 6, 7, 8, float("nan")]])), dim=1
+    ).equal(torch.tensor([10, 26]))
+
+
+def test_torch_nansum_keepdim() -> None:
+    assert torch.nansum(
+        BatchedTensor(torch.tensor([[0, 1, 2, 3, 4], [5, 6, 7, 8, float("nan")]])),
+        dim=1,
+        keepdim=True,
+    ).equal(torch.tensor([[10], [26]]))
+
+
 ################################
 #     Tests for torch.prod     #
 ################################

tests/unit/test_tensor_compatibility.py

@@ -551,7 +551,7 @@ def test_torch_mean(cls: type[BatchedTensor]) -> None:
 @mark.parametrize("dim", (0, 1))
 @mark.parametrize("keepdim", (True, False))
 def test_torch_mean_kwargs(cls: type[BatchedTensor], dim: int, keepdim: bool) -> None:
-    x = torch.rand(6, 10).mul(100)
+    x = torch.rand(6, 10)
     assert torch.mean(cls(x), dim=dim, keepdim=keepdim).equal(
         torch.mean(x, dim=dim, keepdim=keepdim)
     )
@@ -567,7 +567,7 @@ def test_torch_median(cls: type[BatchedTensor]) -> None:
 @mark.parametrize("dim", (0, 1))
 @mark.parametrize("keepdim", (True, False))
 def test_torch_median_kwargs(cls: type[BatchedTensor], dim: int, keepdim: bool) -> None:
-    x = torch.rand(6, 10).mul(100)
+    x = torch.rand(6, 10)
     assert objects_are_equal(
         torch.median(cls(x), dim=dim, keepdim=keepdim), torch.median(x, dim=dim, keepdim=keepdim)
     )
@@ -592,6 +592,23 @@ def test_torch_mul_tensor() -> None:
     )
 
 
+@mark.parametrize("cls", BATCH_CLASSES)
+def test_torch_nansum(cls: type[BatchedTensor]) -> None:
+    x = torch.rand(6, 10)
+    x[x < 10] = float("nan")
+    assert torch.nansum(cls(x)).equal(torch.nansum(x))
+
+
+@mark.parametrize("cls", BATCH_CLASSES)
+@mark.parametrize("dim", (0, 1))
+@mark.parametrize("keepdim", (True, False))
+def test_torch_nansum_kwargs(cls: type[BatchedTensor], dim: int, keepdim: bool) -> None:
+    x = torch.rand(6, 10)
+    assert torch.nansum(cls(x), dim=dim, keepdim=keepdim).equal(
+        torch.nansum(x, dim=dim, keepdim=keepdim)
+    )
+
+
 def test_torch_neg() -> None:
     assert torch.neg(BatchedTensor(torch.full((2, 3), 2.0))).equal(
         BatchedTensor(torch.full((2, 3), -2.0))
@@ -608,7 +625,7 @@ def test_torch_prod(cls: type[BatchedTensor]) -> None:
 @mark.parametrize("dim", (0, 1))
 @mark.parametrize("keepdim", (True, False))
 def test_torch_prod_kwargs(cls: type[BatchedTensor], dim: int, keepdim: bool) -> None:
-    x = torch.rand(6, 10).mul(100)
+    x = torch.rand(6, 10)
     assert torch.prod(cls(x), dim=dim, keepdim=keepdim).equal(
         torch.prod(x, dim=dim, keepdim=keepdim)
     )
@@ -645,5 +662,5 @@ def test_torch_sum(cls: type[BatchedTensor]) -> None:
 @mark.parametrize("dim", (0, 1))
 @mark.parametrize("keepdim", (True, False))
 def test_torch_sum_kwargs(cls: type[BatchedTensor], dim: int, keepdim: bool) -> None:
-    x = torch.rand(6, 10).mul(100)
+    x = torch.rand(6, 10)
     assert torch.sum(cls(x), dim=dim, keepdim=keepdim).equal(torch.sum(x, dim=dim, keepdim=keepdim))

tests/unit/test_tensorseq.py

@@ -7135,6 +7135,31 @@ def test_torch_median_keepdim() -> None:
     )
 
 
+##################################
+#     Tests for torch.nansum     #
+##################################
+
+
+def test_torch_nansum() -> None:
+    assert torch.nansum(
+        BatchedTensorSeq(torch.tensor([[0, 1, 2, 3, 4], [5, 6, 7, 8, float("nan")]]))
+    ).equal(torch.tensor(36))
+
+
+def test_torch_nansum_dim_1() -> None:
+    assert torch.nansum(
+        BatchedTensorSeq(torch.tensor([[0, 1, 2, 3, 4], [5, 6, 7, 8, float("nan")]])), dim=1
+    ).equal(torch.tensor([10, 26]))
+
+
+def test_torch_nansum_keepdim() -> None:
+    assert torch.nansum(
+        BatchedTensorSeq(torch.tensor([[0, 1, 2, 3, 4], [5, 6, 7, 8, float("nan")]])),
+        dim=1,
+        keepdim=True,
+    ).equal(torch.tensor([[10], [26]]))
+
+
 ################################
 #     Tests for torch.prod     #
 ################################