add tests for minor functions

Author: mawright

pytorch_sparse_utils/misc.py

@@ -1,6 +1,7 @@
 import math
 import re
 from typing import TYPE_CHECKING, overload
+import warnings
 
 import torch
 from torch import Tensor
@@ -95,9 +96,12 @@ def unpack_sparse_tensors(batch: dict[str, Tensor]) -> dict[str, Tensor]:
         sparse torch.Tensor format
     """
     if _pytorch_atleast_2_5:
-        raise DeprecationWarning(
-            "`unpack_sparse_tensors` is no longer needed as of Pytorch 2.5",
-            "which added native support for pinned_memory=True for sparse tensors",
+        warnings.warn(
+            (
+                "`unpack_sparse_tensors` is no longer needed as of Pytorch 2.5"
+                "which added native support for pinned_memory=True for sparse tensors"
+            ),
+            DeprecationWarning,
         )
     prefixes_indices = [
         match[0]

tests/test_misc.py

@@ -1,7 +1,16 @@
 import pytest
 import torch
+from torch import Tensor
+from contextlib import ExitStack
+
+from pytorch_sparse_utils.misc import prod, unpack_sparse_tensors, _pytorch_atleast_2_5
+from pytorch_sparse_utils.validation import (
+    validate_atleast_nd,
+    validate_dim_size,
+    validate_nd,
+)
+from . import random_sparse_tensor
 
-from pytorch_sparse_utils.validation import validate_nd, validate_dim_size, validate_atleast_nd
 
 @pytest.mark.cpu_and_cuda
 class TestValidate:
@@ -10,7 +19,7 @@ def test_validate_nd(self, device):
         validate_nd(tensor, 3)
         with pytest.raises(
             (ValueError, torch.jit.Error),  # pyright: ignore[reportArgumentType]
-            match="Expected tensor to be 4D"
+            match="Expected tensor to be 4D",
         ):
             validate_nd(tensor, 4)
 
@@ -19,7 +28,7 @@ def test_validate_at_least_nd(self, device):
         validate_atleast_nd(tensor, 3)
         with pytest.raises(
             (ValueError, torch.jit.Error),  # pyright: ignore[reportArgumentType]
-            match="Expected tensor to have at least"
+            match="Expected tensor to have at least",
         ):
             validate_atleast_nd(tensor, 4)
 
@@ -28,6 +37,43 @@ def test_validate_dim_size(self, device):
         validate_dim_size(tensor, dim=0, expected_size=3)
         with pytest.raises(
             (ValueError, torch.jit.Error),  # pyright: ignore[reportArgumentType]
-            match=r"Expected tensor to have shape\[0\]=4"
+            match=r"Expected tensor to have shape\[0\]=4",
         ):
             validate_dim_size(tensor, dim=0, expected_size=4)
+
+
+@pytest.mark.cpu_and_cuda
+def test_prod(device):
+    test_list = [1, 2, 3]
+    result_list = prod(test_list)
+    assert result_list == 6
+    assert isinstance(result_list, int)
+
+    test_tensor = torch.tensor(test_list, device=device)
+    result_tensor = prod(test_tensor)
+    assert isinstance(result_tensor, Tensor)
+    assert result_tensor == 6
+
+
+@pytest.mark.cpu_and_cuda
+def test_unpack_sparse_tensors(device):
+    sparse_tensor = random_sparse_tensor(
+        [8, 16, 16], [4, 32], 0.5, seed=0, device=device
+    ).coalesce()
+
+    batch_dict = {
+        "X_indices": sparse_tensor.indices(),
+        "X_values": sparse_tensor.values(),
+        "X_shape": torch.tensor(sparse_tensor.shape, device=device),
+    }
+
+    if _pytorch_atleast_2_5:
+        with pytest.warns(DeprecationWarning, match="is no longer needed"):
+            unpacked = unpack_sparse_tensors(batch_dict)
+    else:
+        unpacked = unpack_sparse_tensors(batch_dict)
+
+    X = unpacked["X"]
+    assert torch.equal(X.indices(), sparse_tensor.indices())
+    assert torch.equal(X.values(), sparse_tensor.values())
+    assert X.shape == sparse_tensor.shape