update tests

Author: mawright

pytorch_sparse_utils/conversion.py

@@ -63,11 +63,9 @@ def torch_sparse_to_minkowski(tensor: Tensor):
     assert isinstance(tensor, Tensor)
     assert tensor.is_sparse
     features = tensor.values()
-    coordinates = tensor.indices()
     if features.ndim == 1:
         features = features.unsqueeze(-1)
-        coordinates = coordinates[:-1]
-    coordinates = coordinates.transpose(0, 1).contiguous().int()
+    coordinates = tensor.indices().T.int().contiguous()
     return ME.SparseTensor(
         features, coordinates, requires_grad=tensor.requires_grad, device=tensor.device
     )
@@ -77,11 +75,21 @@ def torch_sparse_to_minkowski(tensor: Tensor):
 def minkowski_to_torch_sparse(
     tensor: Union[Tensor, ME.SparseTensor],
     full_scale_spatial_shape: Optional[Union[Tensor, list[int]]] = None,
+    squeeze: bool = False
 ) -> Tensor:
     """Converts a MinkowskiEngine SparseTensor to an equivalent sparse torch.Tensor
 
     Args:
         tensor (MinkowskiEngine.SparseTensor): Sparse tensor to be converted
+        full_scale_spatial_shape (Optional[Union[list[int], [Tensor]]]): The full
+            extent of the spatial domain on which the sparse data reside.
+            If given, will be used to define the size of the sparse tensor. If not
+            given, the size will be inferred from the indices in the tensor.
+            Default: None
+        squeeze (bool): If True and the feature dimension of the MinkowskiEngine
+            SparseTensor is 1, the returned sparse torch.Tensor will have its values
+            squeezed to 1D shape of [nnz] rather than [nnz, 1]. Raises an error if
+            True and the feature dimension is not 1.
 
     Returns:
         tensor (torch.Tensor): Converted sparse tensor
@@ -102,6 +110,18 @@ def minkowski_to_torch_sparse(
     else:
         max_coords = None
     out = __me_sparse(tensor, min_coords, max_coords)[0].coalesce()
+    if squeeze:
+        if out.values().shape[1] != 1:
+            raise ValueError(
+                "Got `squeeze`=True, but the MinkowskiEngine tensor has a feature "
+                f"dim of {out.values().shape[1]}, not 1."
+            )
+        out = torch.sparse_coo_tensor(
+            out.indices(),
+            out.values().squeeze(-1),
+            out.shape[:-1],
+            is_coalesced=out.is_coalesced()
+        )
     return out
 
 

tests/test_conversion.py

@@ -133,8 +133,8 @@ def test_large_sparse_array(self):
 
 
 @pytest.mark.skipif(not has_minkowskiengine, reason="MinkowskiEngine not installed")
-class TestMinkowskiEngineReal:
-    def test_torch_sparse_to_minkowski_real(self):
+class TestMinkowskiEngine:
+    def test_torch_sparse_to_minkowski(self):
         # Create sparse tensor with 3D coordinates (batch, x, y) and features
         indices = torch.tensor(
             [[0, 0, 1, 1], [0, 1, 0, 1], [0, 0, 1, 1], [3, 4, 5, 6]], dtype=torch.int
@@ -145,12 +145,11 @@ def test_torch_sparse_to_minkowski_real(self):
 
         result = torch_sparse_to_minkowski(tensor)
 
-        assert hasattr(result, "F")  # Features
-        assert hasattr(result, "C")  # Coordinates
+        assert isinstance(result, ME.SparseTensor)
         assert result.F.shape[0] == 4  # Number of points
         assert result.F.shape[1] == 1  # Feature dimension
 
-    def test_minkowski_to_torch_sparse_real(self):
+    def test_minkowski_to_torch_sparse(self):
         # Create a MinkowskiEngine SparseTensor
         coordinates = torch.tensor(
             [[0, 0, 0], [0, 1, 0], [1, 0, 1], [1, 1, 1]], dtype=torch.int
@@ -164,23 +163,53 @@ def test_minkowski_to_torch_sparse_real(self):
         assert result.is_sparse
         assert result.values().numel() == 4
 
-    def test_roundtrip_minkowski(self):
-        # Test roundtrip conversion
+    def test_already_torch_sparse(self):
         indices = torch.tensor(
-            [[0, 0, 1], [0, 1, 0], [0, 1, 1], [2, 3, 4]], dtype=torch.int
+            [[0, 0, 1, 1], [0, 1, 0, 1], [0, 0, 1, 1], [3, 4, 5, 6]], dtype=torch.int
         )
+        values = torch.tensor([1.0, 2.0, 3.0, 4.0])
+        shape = (2, 2, 2, 10)  # batch_size=2, spatial=(2,2), features=10
+        tensor = torch.sparse_coo_tensor(indices, values, shape).coalesce()
+
+        result = minkowski_to_torch_sparse(tensor)
+
+        assert result is tensor
+
+    def test_roundtrip_minkowski(self):
+        # Test roundtrip conversion
+        indices = torch.tensor([[0, 0, 1], [0, 1, 0], [1, 1, 1]], dtype=torch.int).T
         values = torch.tensor([1.0, 2.0, 3.0])
-        shape = (2, 2, 2, 5)
+        shape = (2, 2, 2)
         tensor = torch.sparse_coo_tensor(indices, values, shape).coalesce()
 
         me_tensor = torch_sparse_to_minkowski(tensor)
         back_to_torch = minkowski_to_torch_sparse(
-            me_tensor, full_scale_spatial_shape=[2, 2]
+            me_tensor, full_scale_spatial_shape=[2, 2], squeeze=True
         )
 
-        # Check that we get back similar structure (may not be identical due to coordinate handling)
         assert back_to_torch.is_sparse
-        assert back_to_torch.shape[0] == 2  # batch size preserved
+        assert back_to_torch.shape[0] == 2
+        assert torch.equal(back_to_torch.indices(), tensor.indices())
+        assert torch.equal(back_to_torch.values(), tensor.values())
+
+        # Test with tensor spatial shape
+        back_to_torch_2 = minkowski_to_torch_sparse(
+            me_tensor, full_scale_spatial_shape=torch.tensor([2, 2]), squeeze=True
+        )
+
+        assert torch.equal(back_to_torch.indices(), back_to_torch_2.indices())
+        assert torch.equal(back_to_torch.values(), back_to_torch_2.values())
+
+    def test_squeeze_error(self):
+        # Test trying to squeeze without scalar features
+        coordinates = torch.tensor(
+            [[0, 0, 0], [0, 1, 0], [1, 0, 1], [1, 1, 1]], dtype=torch.int
+        )
+        features = torch.tensor([[1.0, 1.0], [2.0, 2.0], [3.0, 3.0], [4.0, 4.0]])
+        sparse_tensor = ME.SparseTensor(features, coordinates)
+
+        with pytest.raises(ValueError, match="Got `squeeze`=True"):
+            _ = minkowski_to_torch_sparse(sparse_tensor, squeeze=True)
 
 
 @pytest.mark.skipif(not has_spconv, reason="spconv not installed")
@@ -244,7 +273,6 @@ def test_spconv_squeeze_error(self):
         with pytest.raises(ValueError, match="Got `squeeze`=True, but"):
             _ = spconv_to_torch_sparse(sparse_conv_tensor, squeeze=True)
 
-
     def test_roundtrip_spconv(self):
         # Test roundtrip conversion
         indices = torch.tensor([[0, 0, 1], [0, 1, 0], [0, 1, 1], [0, 0, 1]]).T