Repository: cai4cai/torchsparsegradutils (language: python)

Issue:
Add convenience functions for generating random sparse matrices for testing
Rather than repeating code in the unit tests to generate random sparse COO and CSR matrices, it makes more sense to have a utility function to generate them for me, which itself can be unit testing. This should reduce the possibility for errors during unit testing.


Discussion / hints:

Fix patch:
diff --git a/requirements-dev.txt b/requirements-dev.txt
index a4641ad..0bf71d9 100644
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -1,4 +1,5 @@
 torch>=1.13.0
 numpy
 scipy
-jax[cpu]
\ No newline at end of file
+jax[cpu]
+parameterized
\ No newline at end of file
diff --git a/torchsparsegradutils/utils/random_sparse.py b/torchsparsegradutils/utils/random_sparse.py
new file mode 100644
index 0000000..479dd90
--- /dev/null
+++ b/torchsparsegradutils/utils/random_sparse.py
@@ -0,0 +1,273 @@
+"""
+utility functions for generating random sparse matrices
+
+NOTE: sparse COO tensors have indices tensor of size (ndim, nse) and with indices type torch.int64
+NOTE: Sparse CSR The index tensors crow_indices and col_indices should have element type either torch.int64 (default) or torch.int32. 
+      If you want to use MKL-enabled matrix operations, use torch.int32. 
+      This is as a result of the default linking of pytorch being with MKL LP64, which uses 32 bit integer indexing
+NOTE: The batches of sparse CSR tensors are dependent: the number of specified elements in all batches must be the same. 
+      This somewhat  artificial constraint allows efficient storage of the indices of different CSR batches.
+"""
+import warnings
+import torch
+import random
+from torchsparsegradutils.utils.utils import convert_coo_to_csr_indices_values
+
+
+def _gen_indices_2d_coo(nr, nc, nnz, *, dtype=torch.int64, device=torch.device("cpu")):
+    """Generates nnz random unique coordinates in COO format.
+
+    Args:
+        nr (int): Number of rows.
+        nc (int): Number of columns.
+        nnz (int): Number of pairs of coordinates to generate.
+        device (torch.device, optional): Device to generate coordinates on. Defaults to torch.device("cpu").
+
+    Returns:
+        torch.tensor: Tensor of shape [2, nnz] containing the generated coordinates.
+    """
+    assert (
+        nnz <= nr * nc
+    ), "Number of elements (nnz) must be less than or equal to the total number of elements (nr * nc)."
+
+    coordinates = set()
+    while True:
+        r, c = random.randrange(nr), random.randrange(nc)
+        coordinates.add((r, c))
+        if len(coordinates) == nnz:
+            return torch.stack([torch.tensor(co, dtype=dtype, device=device) for co in coordinates], dim=-1)
+
+    # Alternatively, could do:
+    # indices = torch.randperm(nr * nc)[:nnz]
+    # return torch.stack([indices // nc, indices % nc]).to(device)
+
+
+def generate_random_sparse_coo_matrix(
+    size, nnz, *, indices_dtype=torch.int64, values_dtype=torch.float32, device=torch.device("cpu")
+):
+    """Generates a random sparse COO matrix of the specified size and number of non-zero elements.
+
+    Args:
+        size (tuple): Tuple specifying the dimensions of the sparse matrix. The size can be either (num_rows, num_cols) for an unbatched matrix or (batch_size, num_rows, num_cols) for a batched matrix.
+        nnz (int): Number of non-zero values in sparse matrix (number of sparse elements), per batch element. This must be less than or equal to the total number of elements in the matrix of each batch element.
+        indices_dtype (torch.dtype, optional): Data type for indices of sparse tensor. Defaults to torch.int64.
+        values_dtype (torch.dtype, optional): Data type for values of sparse tensor. Defaults to torch.float32.
+        device (torch.device, optional): Device to generate the tensor on. Defaults to torch.device("cpu").
+
+    Raises:
+        ValueError: Raised if size has less than 2 dimensions.
+        ValueError: Raised if size has more than 3 dimensions, as this implementation only supports 1 batch dimension.
+        ValueError: Raised if nnz is greater than the total number of elements (size[-2] * size[-3]).
+        ValueError: Raised if indices_dtype is not torch.int64 for sparse COO tensors.
+
+    Returns:
+        torch.Tensor: Returns a sparse COO tensor of shape size with nnz non-zero elements.
+    """
+
+    if len(size) < 2:
+        raise ValueError("size must have at least 2 dimensions")
+    elif len(size) > 3:
+        raise ValueError("size must have at most 3 dimensions, as this implementation only supports 1 batch dimension")
+
+    if nnz > size[-2] * size[-1]:
+        raise ValueError("nnz must be less than or equal to nr * nc")
+
+    if indices_dtype != torch.int64:
+        raise ValueError("indices_dtype must be torch.int64 for sparse COO tensors")
+
+    if len(size) == 2:
+        coo_indices = _gen_indices_2d_coo(size[-2], size[-1], nnz, dtype=indices_dtype, device=device)
+        values = torch.rand(nnz, dtype=values_dtype, device=device)
+    else:
+        sparse_dim_indices = torch.cat(
+            [_gen_indices_2d_coo(size[-2], size[-1], nnz, dtype=indices_dtype, device=device) for _ in range(size[0])],
+            dim=-1,
+        )
+        batch_dim_indices = (
+            torch.arange(size[0], dtype=indices_dtype, device=device).repeat_interleave(nnz).unsqueeze(0)
+        )
+        coo_indices = torch.cat([batch_dim_indices, sparse_dim_indices])
+        values = torch.rand(nnz * size[0], dtype=values_dtype, device=device)
+
+    return torch.sparse_coo_tensor(coo_indices, values, size, device=device).coalesce()
+
+
+def generate_random_sparse_csr_matrix(
+    size, nnz, *, indices_dtype=torch.int64, values_dtype=torch.float32, device=torch.device("cpu")
+):
+    """Generates a random sparse CSR matrix of the specified size and number of non-zero elements.
+
+    Args:
+        size (tuple): Tuple specifying the dimensions of the sparse matrix. The size can be either (num_rows, num_cols) for an unbatched matrix or (batch_size, num_rows, num_cols) for a batched matrix.
+        nnz (int): Number of non-zero values in sparse matrix (number of sparse elements), per batch element. This must be less than or equal to the total number of elements in the matrix.
+        indices_dtype (torch.dtype, optional): Data type for indices of sparse tensor. Defaults to torch.int64.
+        values_dtype (torch.dtype, optional): Data type for values of sparse tensor. Defaults to torch.float32.
+        device (torch.device, optional): Device to generate the tensor on. Defaults to torch.device("cpu").
+
+    Raises:
+        ValueError: Raised if size has less than 2 dimensions.
+        ValueError: Raised if size has more than 3 dimensions, as this implementation only supports 1 batch dimension.
+        ValueError: Raised if nnz is greater than the total number of elements (size[-2] * size[-3]).
+        ValueError: Raised if indices_dtype is not torch.int64 or torch.int32, as these are the only indices dtypes supported for sparse CSR tensors.
+
+    Returns:
+        torch.Tensor: Returns a sparse CSR tensor of shape size with nnz non-zero elements.
+    """
+    if len(size) < 2:
+        raise ValueError("size must have at least 2 dimensions")
+    elif len(size) > 3:
+        raise ValueError("size must have at most 3 dimensions, as this implementation only supports 1 batch dimension")
+
+    if nnz > size[-2] * size[-1]:
+        raise ValueError("nnz must be less than or equal to nr * nc")
+
+    if (indices_dtype != torch.int64) and (indices_dtype != torch.int32):
+        raise ValueError("indices_dtype must be torch.int64 or torch.int32 for sparse CSR tensors")
+
+    coo_indices = _gen_indices_2d_coo(size[-2], size[-1], nnz, dtype=indices_dtype, device=device)
+    crow_indices, col_indices, _ = convert_coo_to_csr_indices_values(coo_indices, size[-2], values=None)
+
+    if len(size) == 2:
+        values = torch.rand(nnz, dtype=values_dtype, device=device)
+    else:
+        crow_indices = crow_indices.repeat(size[0], 1)
+        col_indices = col_indices.repeat(size[0], 1)
+        values = torch.rand((size[0], nnz), dtype=values_dtype, device=device)
+
+    return torch.sparse_csr_tensor(crow_indices, col_indices, values, size, device=device)
+
+
+# Square strictly Triangular:
+
+
+def _gen_indices_2d_coo_strictly_tri(n, nnz, *, upper=True, dtype=torch.int64, device=torch.device("cpu")):
+    """Generates nnz random unique COO coordinates for a square matrix with either strictly lower triangular or strictly upper triangular coordinates.
+
+    Args:
+        n (int): Size of the square matrix (number of rows and columns).
+        nnz (int): Number of elements the generated COO indices represent.
+        upper (bool, optional): Flag indicating whether to generate strictly upper triangular coordinates, or strictly lower triangular coordinates. Defaults to True.
+        dtype (torch.dtype, optional): Data type of the tensors. Defaults to torch.int64.
+        device (torch.device, optional): Device to generate coordinates on. Defaults to torch.device("cpu").
+
+    Returns:
+        torch.Tensor: Tensor of shape [2, nnz] containing the generated COO indices.
+    """
+    assert (
+        nnz <= n * (n - 1) // 2
+    ), "Number of elements (nnz) must be less than or equal to the total number of elements (n * (n - 1) // 2)."
+
+    coordinates = set()
+    while True:
+        r, c = random.randrange(n), random.randrange(n)
+        if (r < c and upper) or (r > c and not upper):
+            coordinates.add((r, c))
+        if len(coordinates) == nnz:
+            return torch.stack([torch.tensor(co, dtype=dtype, device=device) for co in coordinates], dim=-1).to(device)
+
+
+def generate_random_sparse_strictly_triangular_coo_matrix(
+    size, nnz, *, upper=True, indices_dtype=torch.int64, values_dtype=torch.float32, device=torch.device("cpu")
+):
+    """Generates a random sparse COO square matrix with strictly upper or lower triangular coordinates.
+
+    Args:
+        size (tuple): Tuple specifying the dimensions of the sparse matrix. The size can be either (num_rows, num_cols) for an unbatched matrix or (batch_size, num_rows, num_cols) for a batched matrix. The number of rows and columns must be equal.
+        nnz (int): Number of non-zero values in sparse matrix (number of sparse elements), per batch element. nnz must be less than or equal to (n * n-1)/2, where n is the number of rows or columns.
+        upper (bool, optional): If True, generates strictly upper triangular indices. If False, generates strictly lower triangular indices. Defaults to True.
+        indices_dtype (torch.dtype, optional): Data type for indices of sparse tensor. Defaults to torch.int64.
+        values_dtype (torch.dtype, optional): Data type for values of sparse tensor. Defaults to torch.float32.
+        device (torch.device, optional): Device to generate the coordinates on. Defaults to torch.device("cpu").
+
+    Raises:
+        ValueError: Raised if size has less than 2 dimensions.
+        ValueError: Raised if size has more than 3 dimensions, as this implementation only supports 1 batch dimension.
+        ValueError: Raised if size is not a square matrix (n, n) or batched square matrix (b, n, n).
+        ValueError: Raised if nnz is greater than (n * n-1)/2, where n is the number of rows or columns.
+        ValueError: Raised if indices_dtype is not torch.int64 for sparse COO tensors.
+
+    Returns:
+        torch.Tensor: Returns a square strictly upper or lower sparse COO tensor of shape size with nnz non-zero elements.
+    """
+    if len(size) < 2:
+        raise ValueError("size must have at least 2 dimensions")
+    elif len(size) > 3:
+        raise ValueError("size must have at most 3 dimensions, as this implementation only supports 1 batch dimension")
+
+    if size[-2] != size[-1]:
+        raise ValueError("size must be a square matrix (n, n) or batched square matrix (b, n, n)")
+
+    if nnz > size[-2] * (size[-2] - 1) // 2:
+        raise ValueError("nnz must be less than or equal to (n * n-1)/2, where n is the number of rows or columns")
+
+    if indices_dtype != torch.int64:
+        raise ValueError("indices_dtype must be torch.int64 for sparse COO tensors")
+
+    if len(size) == 2:
+        coo_indices = _gen_indices_2d_coo_strictly_tri(size[-2], nnz, upper=upper, dtype=indices_dtype, device=device)
+        values = torch.rand(nnz, dtype=values_dtype, device=device)
+    else:
+        sparse_dim_indices = torch.cat(
+            [
+                _gen_indices_2d_coo_strictly_tri(size[-2], nnz, upper=upper, dtype=indices_dtype, device=device)
+                for _ in range(size[0])
+            ],
+            dim=-1,
+        )
+        batch_dim_indices = batch_dim_indices = (
+            torch.arange(size[0], dtype=indices_dtype, device=device).repeat_interleave(nnz).unsqueeze(0)
+        )
+        coo_indices = torch.cat([batch_dim_indices, sparse_dim_indices])
+        values = torch.rand(nnz * size[0], dtype=values_dtype, device=device)
+
+    return torch.sparse_coo_tensor(coo_indices, values, size, device=device).coalesce()
+
+
+def generate_random_sparse_strictly_triangular_csr_matrix(
+    size, nnz, *, upper=True, indices_dtype=torch.int64, values_dtype=torch.float32, device=torch.device("cpu")
+):
+    """Generates a random sparse CSR square matrix with strictly upper or lower triangular coordinates.
+
+    Args:
+        size (tuple): Tuple specifying the dimensions of the sparse matrix. The size can be either (num_rows, num_cols) for an unbatched matrix or (batch_size, num_rows, num_cols) for a batched matrix. The number of rows and columns must be equal.
+        nnz (int): Number of non-zero values in sparse matrix (number of sparse elements), per batch element. nnz must be less than or equal to (n * n-1)/2, where n is the number of rows or columns.
+        upper (bool, optional): If True, generates strictly upper triangular indices. If False, generates strictly lower triangular indices. Defaults to True.
+        indices_dtype (torch.dtype, optional): Data type for indices of sparse tensor. Defaults to torch.int64.
+        values_dtype (torch.dtype, optional): Data type for values of sparse tensor. Defaults to torch.float32.
+        device (torch.device, optional): Device to generate the coordinates on. Defaults to torch.device("cpu").
+
+    Raises:
+        ValueError: Raised if size has less than 2 dimensions.
+        ValueError: Raised if size has more than 3 dimensions, as this implementation only supports 1 batch dimension.
+        ValueError: Raised if size is not a square matrix (n, n) or batched square matrix (b, n, n).
+        ValueError: Raised if nnz is greater than (n * n-1)/2, where n is the number of rows or columns.
+        ValueError: Raised if indices_dtype is not torch.int64 or torch.int32, as these are the only indices dtypes supported for sparse CSR tensors.
+
+    Returns:
+        torch.Tensor: Returns a square strictly upper or lower sparse CSR tensor of shape size with nnz non-zero elements.
+    """
+    if len(size) < 2:
+        raise ValueError("size must have at least 2 dimensions")
+    elif len(size) > 3:
+        raise ValueError("size must have at most 3 dimensions, as this implementation only supports 1 batch dimension")
+
+    if size[-2] != size[-1]:
+        raise ValueError("size must be a square matrix (n, n) or batched square matrix (b, n, n)")
+
+    if nnz > size[-2] * (size[-2] - 1) // 2:
+        raise ValueError("nnz must be less than or equal to (n * n-1)/2, where n is the number of rows or columns")
+
+    if (indices_dtype != torch.int64) and (indices_dtype != torch.int32):
+        raise ValueError("indices_dtype must be torch.int64 or torch.int32 for sparse CSR tensors")
+
+    coo_indices = _gen_indices_2d_coo_strictly_tri(size[-2], nnz, upper=upper, dtype=indices_dtype, device=device)
+    crow_indices, col_indices, _ = convert_coo_to_csr_indices_values(coo_indices, size[-2], values=None)
+
+    if len(size) == 2:
+        values = torch.rand(nnz, dtype=values_dtype, device=device)
+    else:
+        crow_indices = crow_indices.repeat(size[0], 1)
+        col_indices = col_indices.repeat(size[0], 1)
+        values = torch.rand((size[0], nnz), dtype=values_dtype, device=device)
+
+    return torch.sparse_csr_tensor(crow_indices, col_indices, values, size, device=device)
diff --git a/torchsparsegradutils/utils/utils.py b/torchsparsegradutils/utils/utils.py
new file mode 100644
index 0000000..3bc475c
--- /dev/null
+++ b/torchsparsegradutils/utils/utils.py
@@ -0,0 +1,145 @@
+import torch
+
+
+def _sort_coo_indices(indices):
+    """
+    Sorts COO (Coordinate List Format) indices in ascending order and returns a permutation tensor that indicates the indices in the original data that result in a sorted tensor.
+
+    This function can support both unbatched and batched COO indices, and essentially performs the same operation as .coalesce() called on a COO tensor.
+    The advantage is that COO coordinates can be sorted prior to conversion to CSR, without having to use the torch.sparse_coo_tensor object which only supports int64 indices.
+
+    Args:
+        indices (torch.Tensor): The input indices in COO format to be sorted.
+
+    Returns:
+        torch.Tensor: A tensor containing sorted indices.
+        torch.Tensor: A permutation tensor that contains the indices in the original tensor that give the sorted tensor.
+    """
+    indices_sorted, permutation = torch.unique(indices, dim=-1, sorted=True, return_inverse=True)
+    return indices_sorted.contiguous(), torch.argsort(permutation)
+
+
+def _compress_row_indices(row_indices, num_rows):
+    """Compresses COO row indices to CSR crow indices.
+
+
+    Args:
+        row_indices (torch.Tensor): Tensor of COO row indices.
+        num_rows (int): Number of rows in the matrix.
+
+
+    Returns:
+        torch.Tensor: Compressed CSR crow indices.
+    """
+    # Compute the number of non-zero elements in each row
+    counts = torch.bincount(row_indices, minlength=num_rows).to(row_indices.dtype)
+
+    # Compute the cumulative sum of counts to get CSR indices
+    crow_indices = torch.cat([torch.zeros(1, dtype=row_indices.dtype, device=counts.device), counts.cumsum_(dim=0)])
+
+    return crow_indices
+
+
+def convert_coo_to_csr_indices_values(coo_indices, num_rows, values=None):
+    """Converts COO row and column indices to CSR crow and col indices.
+    Supports batched indices, which would have shape [3, nnz]. Or, [2, nnz] for unbatched indices.
+    This function sorts the row and column indices (similar to torch.sparse_coo_tensor.coalesce()) and then compresses the row indices to CSR crow indices.
+    If values are provided, the values tensor is permuted according to the sorted COO indices.
+    If no values are provided, the permutation indices are returned.
+
+
+    Args:
+        coo_indices (torch.Tensor): Tensor of COO indices
+        num_rows (int): Number of rows in the matrix.
+
+
+    Returns:
+        torch.Tensor: Compressed CSR crow indices.
+        torch.Tensor: CSR Col indices.
+        torch.Tensor: Permutation indices from sorting COO indices. Or permuted values if values are provided.
+    """
+    if coo_indices.shape[0] < 2:
+        raise ValueError(
+            f"Indices tensor must have at least 2 rows (row and column indices). Got {coo_indices.shape[0]} rows."
+        )
+    elif coo_indices.shape[0] > 3:
+        raise ValueError(
+            f"Current implementation only supports single batch diomension, therefore indices tensor must have at most 3 rows (batch, row and column indices). Got {coo_indices.shape[0]} rows."
+        )
+
+    if coo_indices[-2].max() >= num_rows:
+        raise ValueError(
+            f"Row indices must be less than num_rows ({num_rows}). Got max row index {coo_indices[-2].max()}"
+        )
+
+    if values != None and values.shape[0] != coo_indices.shape[1]:
+        raise ValueError(
+            f"Number of values ({values.shape[0]}) does not match number of indices ({coo_indices.shape[1]})"
+        )
+
+    coo_indices, permutation = _sort_coo_indices(coo_indices)
+
+    if coo_indices.shape[0] == 2:
+        row_indices, col_indices = coo_indices
+        crow_indices = _compress_row_indices(row_indices, num_rows)
+
+        values = values[permutation] if values is not None else permutation
+
+    else:
+        batch_indices, row_indices, col_indices = coo_indices
+        crow_indices = torch.cat(
+            [
+                _compress_row_indices(row_indices[batch_indices == batch], num_rows)
+                for batch in torch.unique(batch_indices)
+            ]
+        )
+        num_batches = torch.unique(batch_indices).shape[0]
+
+        crow_indices = crow_indices.reshape(num_batches, -1)
+        col_indices = col_indices.reshape(num_batches, -1)
+
+        values = values[permutation] if values is not None else permutation
+
+        values = values.reshape(num_batches, -1)
+
+    return crow_indices, col_indices, values
+
+
+def convert_coo_to_csr(sparse_coo_tensor):
+    """Converts a COO sparse tensor to CSR format. COO tensor can have optional single leading batch dimension.
+
+    Args:
+        sparse_coo_tensor (torch.Tensor): COO sparse tensor.
+
+
+    Returns:
+        torch.Tensor: CSR sparse tensor.
+    """
+    if sparse_coo_tensor.layout == torch.sparse_coo:
+        crow_indices, col_indices, values = convert_coo_to_csr_indices_values(
+            sparse_coo_tensor.indices(), sparse_coo_tensor.size()[-2], sparse_coo_tensor.values()
+        )
+        return torch.sparse_csr_tensor(crow_indices, col_indices, values, sparse_coo_tensor.size())
+    else:
+        raise ValueError(f"Unsupported layout: {sparse_coo_tensor.layout}")
+
+
+def _demcompress_crow_indices(crow_indices, num_rows):
+    """Decompresses CSR crow indices to COO row indices.
+
+
+    Args:
+        csr_crow_indices (torch.Tensor): Tensor of CSR crow indices.
+        num_rows (int): Number of rows in the matrix.
+
+
+    Returns:
+        torch.Tensor: Decompressed COO row indices.
+    """
+
+    row_indices = torch.repeat_interleave(
+        torch.arange(num_rows, dtype=crow_indices.dtype, device=crow_indices.device),
+        crow_indices[1:] - crow_indices[:-1],
+    )
+
+    return row_indices


Test patch:
diff --git a/tests/test_random.py b/tests/test_random.py
new file mode 100644
index 0000000..b8a0f99
--- /dev/null
+++ b/tests/test_random.py
@@ -0,0 +1,475 @@
+import unittest
+from parameterized import parameterized, parameterized_class
+import torch
+from torchsparsegradutils.utils.random_sparse import (
+    generate_random_sparse_coo_matrix,
+    generate_random_sparse_csr_matrix,
+    generate_random_sparse_strictly_triangular_coo_matrix,
+    generate_random_sparse_strictly_triangular_csr_matrix,
+)
+
+# https://pytorch.org/docs/stable/generated/torch.sparse.check_sparse_tensor_invariants.html#torch.sparse.check_sparse_tensor_invariants
+torch.sparse.check_sparse_tensor_invariants.enable()
+
+
+@parameterized_class(
+    (
+        "name",
+        "device",
+    ),
+    [
+        ("CPU", torch.device("cpu")),
+        (
+            "CUDA",
+            torch.device("cuda"),
+        ),
+    ],
+)
+class TestGenRandomCOO(unittest.TestCase):
+    def setUp(self) -> None:
+        if not torch.cuda.is_available() and self.device == torch.device("cuda"):
+            self.skipTest(f"Skipping {self.__class__.__name__} since CUDA is not available")
+
+    # error handling:
+    def test_too_few_dims(self):
+        with self.assertRaises(ValueError):
+            generate_random_sparse_coo_matrix(torch.Size([16]), 6)
+
+    def test_too_many_dims(self):
+        with self.assertRaises(ValueError):
+            generate_random_sparse_coo_matrix(torch.Size([4, 4, 8, 8]), 32)
+
+    def test_too_many_nnz(self):
+        with self.assertRaises(ValueError):
+            generate_random_sparse_coo_matrix(torch.Size([4, 4]), 17)
+
+    @parameterized.expand(
+        [
+            ("int08", torch.int8),
+            ("int16", torch.int16),
+            ("int32", torch.int32),
+        ]
+    )
+    def test_incompatible_indices_dtype(self, _, indices_dtype):
+        with self.assertRaises(ValueError):
+            generate_random_sparse_coo_matrix(torch.Size([4, 4]), 12, indices_dtype=indices_dtype)
+
+    # basic properties:
+    def test_device(self):
+        A = generate_random_sparse_coo_matrix(torch.Size([4, 4]), 12, device=self.device)
+        self.assertEqual(A.device.type, self.device.type)
+
+    @parameterized.expand(
+        [
+            ("int64", torch.int64),
+        ]
+    )  # NOTE: only torch.int64 is supported for COO indices, all other dtypes are casted to torch.int64
+    def test_indices_dtype(self, _, indices_dtype):
+        A = generate_random_sparse_coo_matrix(torch.Size([4, 4]), 12, indices_dtype=indices_dtype, device=self.device)
+        self.assertEqual(A.indices().dtype, indices_dtype)
+
+    @parameterized.expand(
+        [
+            ("float16", torch.float16),
+            ("float32", torch.float32),
+            ("float64", torch.float64),
+        ]
+    )
+    def test_values_dtype(self, _, values_dtype):
+        A = generate_random_sparse_coo_matrix(torch.Size([4, 4]), 12, values_dtype=values_dtype, device=self.device)
+        self.assertEqual(A.values().dtype, values_dtype)
+
+    @parameterized.expand(
+        [
+            ("4x4", torch.Size([4, 4]), 12),
+            ("2x4x4", torch.Size([2, 4, 4]), 12),
+            ("8x16", torch.Size([8, 16]), 32),
+            ("4x8x16", torch.Size([4, 8, 16]), 32),
+        ]
+    )
+    def test_size(self, _, size, nnz):
+        A = generate_random_sparse_coo_matrix(size, nnz, device=self.device)
+        self.assertEqual(A.size(), size)
+
+    @parameterized.expand(
+        [
+            ("4x4", torch.Size([4, 4]), 12),
+            ("2x4x4", torch.Size([2, 4, 4]), 12),
+            ("8x16", torch.Size([8, 16]), 32),
+            ("4x8x16", torch.Size([4, 8, 16]), 32),
+        ]
+    )
+    def test_nnz(self, _, size, nnz):
+        # NOTE: COO tensors return ._nnz() over all batch elements
+        A = generate_random_sparse_coo_matrix(size, nnz, device=self.device)
+        self.assertEqual(A._nnz(), nnz * ((1,) + size)[-3])
+
+
+@parameterized_class(
+    (
+        "name",
+        "device",
+    ),
+    [
+        ("CPU", torch.device("cpu")),
+        (
+            "CUDA",
+            torch.device("cuda"),
+        ),
+    ],
+)
+class TestGenRandomCSR(unittest.TestCase):
+    def setUp(self) -> None:
+        if not torch.cuda.is_available() and self.device == torch.device("cuda"):
+            self.skipTest(f"Skipping {self.__class__.__name__} since CUDA is not available")
+
+    # error handling:
+    def test_too_few_dims(self):
+        with self.assertRaises(ValueError):
+            generate_random_sparse_csr_matrix(torch.Size([16]), 6)
+
+    def test_too_many_dims(self):
+        with self.assertRaises(ValueError):
+            generate_random_sparse_csr_matrix(torch.Size([4, 4, 8, 8]), 32)
+
+    def test_too_many_nnz(self):
+        with self.assertRaises(ValueError):
+            generate_random_sparse_csr_matrix(torch.Size([4, 4]), 17)
+
+    @parameterized.expand(
+        [
+            ("int08", torch.int8),
+            ("int16", torch.int16),
+        ]
+    )
+    def test_incompatible_indices_dtype(self, _, indices_dtype):
+        with self.assertRaises(ValueError):
+            generate_random_sparse_coo_matrix(torch.Size([4, 4]), 12, indices_dtype=indices_dtype)
+
+    # basic properties:
+    def test_device(self):
+        A = generate_random_sparse_csr_matrix(torch.Size([4, 4]), 12, device=self.device)
+        self.assertEqual(A.device.type, self.device.type)
+
+    @parameterized.expand(
+        [
+            ("int32", torch.int32),
+            ("int64", torch.int64),
+        ]
+    )  # NOTE: Only int32 and int64 are supported for CSR indices
+    def test_indices_dtype(self, _, indices_dtype):
+        A = generate_random_sparse_csr_matrix(torch.Size([4, 4]), 12, indices_dtype=indices_dtype, device=self.device)
+        self.assertEqual(A.crow_indices().dtype, indices_dtype)
+        self.assertEqual(A.col_indices().dtype, indices_dtype)
+
+    @parameterized.expand(
+        [
+            ("float16", torch.float16),
+            ("float32", torch.float32),
+            ("float64", torch.float64),
+        ]
+    )
+    def test_values_dtype(self, _, values_dtype):
+        A = generate_random_sparse_csr_matrix(torch.Size([4, 4]), 12, values_dtype=values_dtype, device=self.device)
+        self.assertEqual(A.values().dtype, values_dtype)
+
+    @parameterized.expand(
+        [
+            ("4x4", torch.Size([4, 4]), 12),
+            ("2x4x4", torch.Size([2, 4, 4]), 12),
+            ("8x16", torch.Size([8, 16]), 32),
+            ("4x8x16", torch.Size([4, 8, 16]), 32),
+        ]
+    )
+    def test_size(self, _, size, nnz):
+        A = generate_random_sparse_csr_matrix(size, nnz, device=self.device)
+        self.assertEqual(A.size(), size)
+
+    @parameterized.expand(
+        [
+            ("4x4", torch.Size([4, 4]), 12),
+            ("2x4x4", torch.Size([2, 4, 4]), 12),
+            ("8x16", torch.Size([8, 16]), 32),
+            ("4x8x16", torch.Size([4, 8, 16]), 32),
+        ]
+    )
+    def test_nnz(self, _, size, nnz):
+        A = generate_random_sparse_csr_matrix(size, nnz, device=self.device)
+        # NOTE: CSR tensors return ._nnz() per batch element
+        self.assertEqual(A._nnz(), nnz)
+
+
+# Strictly triangular tests:
+
+
+@parameterized_class(
+    (
+        "name",
+        "device",
+    ),
+    [
+        ("Lower_CPU", torch.device("cpu")),
+        ("CUDA", torch.device("cuda")),
+    ],
+)
+class TestGenRandomStrictlyTriCOO(unittest.TestCase):
+    def setUp(self) -> None:
+        if not torch.cuda.is_available() and self.device == torch.device("cuda"):
+            self.skipTest(f"Skipping {self.__class__.__name__} since CUDA is not available")
+
+    # error handling:
+    def test_too_few_dims(self):
+        with self.assertRaises(ValueError):
+            generate_random_sparse_strictly_triangular_coo_matrix(torch.Size([16]), 6)
+
+    def test_too_many_dims(self):
+        with self.assertRaises(ValueError):
+            generate_random_sparse_strictly_triangular_coo_matrix(torch.Size([4, 4, 8, 8]), 12)
+
+    def test_too_many_nnz(self):
+        with self.assertRaises(ValueError):
+            limit = 4 * (4 - 1) // 2
+            generate_random_sparse_strictly_triangular_coo_matrix(torch.Size([4, 4]), limit + 1)
+
+    @parameterized.expand(
+        [
+            ("int08", torch.int8),
+            ("int16", torch.int16),
+            ("int32", torch.int32),
+        ]
+    )
+    def test_incompatible_indices_dtype(self, _, indices_dtype):
+        with self.assertRaises(ValueError):
+            generate_random_sparse_strictly_triangular_coo_matrix(torch.Size([4, 4]), 5, indices_dtype=indices_dtype)
+
+    # basic properties:
+    def test_device(self):
+        A = generate_random_sparse_strictly_triangular_coo_matrix(torch.Size([4, 4]), 5, device=self.device)
+        self.assertEqual(A.device.type, self.device.type)
+
+    @parameterized.expand(
+        [
+            ("int64", torch.int64),
+        ]
+    )  # NOTE: only torch.int64 is supported for COO indices, all other dtypes are casted to torch.int64
+    def test_indices_dtype(self, _, indices_dtype):
+        A = generate_random_sparse_strictly_triangular_coo_matrix(
+            torch.Size([4, 4]), 5, indices_dtype=indices_dtype, device=self.device
+        )
+        self.assertEqual(A.indices().dtype, indices_dtype)
+
+    @parameterized.expand(
+        [
+            ("float16", torch.float16),
+            ("float32", torch.float32),
+            ("float64", torch.float64),
+        ]
+    )
+    def test_values_dtype(self, _, values_dtype):
+        A = generate_random_sparse_strictly_triangular_coo_matrix(
+            torch.Size([4, 4]), 5, values_dtype=values_dtype, device=self.device
+        )
+        self.assertEqual(A.values().dtype, values_dtype)
+
+    # specific properties:
+    @parameterized.expand(
+        [
+            ("4x4", torch.Size([4, 4]), 4),
+            ("2x4x4", torch.Size([2, 4, 4]), 4),
+            ("8x8", torch.Size([8, 8]), 28),
+            ("4x8x8", torch.Size([4, 8, 8]), 28),
+        ]
+    )
+    def test_size_upper(self, _, size, nnz):
+        A = generate_random_sparse_strictly_triangular_coo_matrix(size, nnz, upper=True, device=self.device)
+        self.assertEqual(A.size(), size)
+
+    @parameterized.expand(
+        [
+            ("4x4", torch.Size([4, 4]), 4),
+            ("2x4x4", torch.Size([2, 4, 4]), 4),
+            ("8x8", torch.Size([8, 8]), 28),
+            ("4x8x8", torch.Size([4, 8, 8]), 28),
+        ]
+    )
+    def test_size_lower(self, _, size, nnz):
+        A = generate_random_sparse_strictly_triangular_coo_matrix(size, nnz, upper=False, device=self.device)
+        self.assertEqual(A.size(), size)
+
+    @parameterized.expand(
+        [
+            ("4x4_12nnz", torch.Size([4, 4]), 5),
+            ("2x4x4_12nnz", torch.Size([2, 4, 4]), 4),
+            ("8x8_22nnz", torch.Size([8, 8]), 22),
+            ("4x8x8_28nnz", torch.Size([4, 8, 8]), 28),
+        ]
+    )
+    def test_nnz_upper(self, _, size, nnz):
+        # NOTE: COO tensors return ._nnz() over all batch elements
+        A = generate_random_sparse_strictly_triangular_coo_matrix(size, nnz, upper=True, device=self.device)
+        self.assertEqual(A._nnz(), nnz * ((1,) + size)[-3])
+
+    @parameterized.expand(
+        [
+            ("4x4_12nnz", torch.Size([4, 4]), 5),
+            ("2x4x4_12nnz", torch.Size([2, 4, 4]), 4),
+            ("8x8_22nnz", torch.Size([8, 8]), 22),
+            ("4x8x8_28nnz", torch.Size([4, 8, 8]), 28),
+        ]
+    )
+    def test_nnz_lower(self, _, size, nnz):
+        # NOTE: COO tensors return ._nnz() over all batch elements
+        A = generate_random_sparse_strictly_triangular_coo_matrix(size, nnz, upper=False, device=self.device)
+        self.assertEqual(A._nnz(), nnz * ((1,) + size)[-3])
+
+    def test_is_strictly_lower(self):
+        A = generate_random_sparse_strictly_triangular_coo_matrix(
+            torch.Size([4, 4]), 5, upper=False, device=self.device
+        )
+        Ad = A.to_dense()
+        self.assertTrue(torch.equal(Ad, Ad.tril(-1)))
+
+    def test_is_strictly_upper(self):
+        A = generate_random_sparse_strictly_triangular_coo_matrix(
+            torch.Size([4, 4]), 5, upper=False, device=self.device
+        )
+        Ad = A.to_dense()
+        self.assertTrue(torch.equal(Ad, Ad.tril(1)))
+
+
+@parameterized_class(
+    (
+        "name",
+        "device",
+    ),
+    [
+        ("CPU", torch.device("cpu")),
+        (
+            "CUDA",
+            torch.device("cuda"),
+        ),
+    ],
+)
+class TestGenRandomStrictlyTriCSR(unittest.TestCase):
+    def setUp(self) -> None:
+        if not torch.cuda.is_available() and self.device == torch.device("cuda"):
+            self.skipTest(f"Skipping {self.__class__.__name__} since CUDA is not available")
+
+    # error handling:
+    def test_too_few_dims(self):
+        with self.assertRaises(ValueError):
+            generate_random_sparse_strictly_triangular_csr_matrix(torch.Size([16]), 6)
+
+    def test_too_many_dims(self):
+        with self.assertRaises(ValueError):
+            generate_random_sparse_strictly_triangular_csr_matrix(torch.Size([4, 4, 8, 8]), 12)
+
+    def test_too_many_nnz(self):
+        with self.assertRaises(ValueError):
+            limit = 4 * (4 - 1) // 2
+            generate_random_sparse_strictly_triangular_csr_matrix(torch.Size([4, 4]), limit + 1)
+
+    @parameterized.expand(
+        [
+            ("int08", torch.int8),
+            ("int16", torch.int16),
+        ]
+    )
+    def test_incompatible_indices_dtype(self, _, indices_dtype):
+        with self.assertRaises(ValueError):
+            generate_random_sparse_strictly_triangular_csr_matrix(torch.Size([4, 4]), 5, indices_dtype=indices_dtype)
+
+    # basic properties:
+    def test_device(self):
+        A = generate_random_sparse_strictly_triangular_csr_matrix(torch.Size([4, 4]), 5, device=self.device)
+        self.assertEqual(A.device.type, self.device.type)
+
+    @parameterized.expand(
+        [
+            ("int32", torch.int32),
+            ("int64", torch.int64),
+        ]
+    )
+    def test_indices_dtype(self, _, indices_dtype):
+        A = generate_random_sparse_strictly_triangular_csr_matrix(
+            torch.Size([4, 4]), 5, indices_dtype=indices_dtype, device=self.device
+        )
+        self.assertEqual(A.crow_indices().dtype, indices_dtype)
+        self.assertEqual(A.col_indices().dtype, indices_dtype)
+
+    @parameterized.expand(
+        [
+            ("float16", torch.float16),
+            ("float32", torch.float32),
+            ("float64", torch.float64),
+        ]
+    )
+    def test_values_dtype(self, _, values_dtype):
+        A = generate_random_sparse_strictly_triangular_csr_matrix(
+            torch.Size([4, 4]), 5, values_dtype=values_dtype, device=self.device
+        )
+        self.assertEqual(A.values().dtype, values_dtype)
+
+    # specific properties:
+    @parameterized.expand(
+        [
+            ("4x4", torch.Size([4, 4]), 4),
+            ("2x4x4", torch.Size([2, 4, 4]), 4),
+            ("8x8", torch.Size([8, 8]), 28),
+            ("4x8x8", torch.Size([4, 8, 8]), 28),
+        ]
+    )
+    def test_size_upper(self, _, size, nnz):
+        A = generate_random_sparse_strictly_triangular_csr_matrix(size, nnz, upper=True, device=self.device)
+        self.assertEqual(A.size(), size)
+
+    @parameterized.expand(
+        [
+            ("4x4", torch.Size([4, 4]), 4),
+            ("2x4x4", torch.Size([2, 4, 4]), 4),
+            ("8x8", torch.Size([8, 8]), 28),
+            ("4x8x8", torch.Size([4, 8, 8]), 28),
+        ]
+    )
+    def test_size_lower(self, _, size, nnz):
+        A = generate_random_sparse_strictly_triangular_csr_matrix(size, nnz, upper=False, device=self.device)
+        self.assertEqual(A.size(), size)
+
+    @parameterized.expand(
+        [
+            ("4x4_12nnz", torch.Size([4, 4]), 5),
+            ("2x4x4_12nnz", torch.Size([2, 4, 4]), 4),
+            ("8x8_22nnz", torch.Size([8, 8]), 22),
+            ("4x8x8_28nnz", torch.Size([4, 8, 8]), 28),
+        ]
+    )
+    def test_nnz_upper(self, _, size, nnz):
+        A = generate_random_sparse_strictly_triangular_csr_matrix(size, nnz, upper=True, device=self.device)
+        # NOTE: CSR tensors return ._nnz() per batch element
+        self.assertEqual(A._nnz(), nnz)
+
+    @parameterized.expand(
+        [
+            ("4x4_12nnz", torch.Size([4, 4]), 5),
+            ("2x4x4_12nnz", torch.Size([2, 4, 4]), 4),
+            ("8x8_22nnz", torch.Size([8, 8]), 22),
+            ("4x8x8_28nnz", torch.Size([4, 8, 8]), 28),
+        ]
+    )
+    def test_nnz_lower(self, _, size, nnz):
+        A = generate_random_sparse_strictly_triangular_csr_matrix(size, nnz, upper=False, device=self.device)
+        # NOTE: CSR tensors return ._nnz() per batch element
+        self.assertEqual(A._nnz(), nnz)
+
+    def test_is_strictly_upper(self):
+        A = generate_random_sparse_strictly_triangular_csr_matrix(torch.Size([4, 4]), 5, upper=True, device=self.device)
+        Ad = A.to_dense()
+        self.assertTrue(torch.equal(Ad, Ad.triu(1)))
+
+    def test_is_strictly_lower(self):
+        A = generate_random_sparse_strictly_triangular_csr_matrix(
+            torch.Size([4, 4]), 5, upper=False, device=self.device
+        )
+        Ad = A.to_dense()
+        print(Ad)
+        self.assertTrue(torch.equal(Ad, Ad.tril(-1)))
diff --git a/tests/test_utils.py b/tests/test_utils.py
new file mode 100644
index 0000000..0266344
--- /dev/null
+++ b/tests/test_utils.py
@@ -0,0 +1,196 @@
+import torch
+import unittest
+from parameterized import parameterized_class, parameterized
+from torchsparsegradutils.utils.random_sparse import (
+    generate_random_sparse_coo_matrix,
+    generate_random_sparse_csr_matrix,
+)
+
+from torchsparsegradutils.utils.utils import (
+    _compress_row_indices,
+    _demcompress_crow_indices,
+    _sort_coo_indices,
+    convert_coo_to_csr,
+)
+
+# https://pytorch.org/docs/stable/generated/torch.sparse.check_sparse_tensor_invariants.html#torch.sparse.check_sparse_tensor_invariants
+torch.sparse.check_sparse_tensor_invariants.enable()
+
+
+@parameterized_class(
+    (
+        "name",
+        "device",
+    ),
+    [
+        ("CPU", torch.device("cpu")),
+        (
+            "CUDA",
+            torch.device("cuda"),
+        ),
+    ],
+)
+class TestSortCOOIndices(unittest.TestCase):
+    def setUp(self) -> None:
+        if not torch.cuda.is_available() and self.device == torch.device("cuda"):
+            self.skipTest(f"Skipping {self.__class__.__name__} since CUDA is not available")
+
+    def test_unbatched_sort(self):
+        nr, nc = 4, 4
+        indices = torch.randperm(nr * nc, device=self.device)
+        indices = torch.stack([indices // nc, indices % nc])
+
+        values = torch.arange(16, device=self.device)
+        sorted_indices_coalesced = torch.sparse_coo_tensor(indices, values).coalesce().indices()
+        coalesce_permutation = torch.sparse_coo_tensor(indices, values).coalesce().values()
+        sorted_indices, permutation = _sort_coo_indices(indices)
+        self.assertTrue(torch.equal(sorted_indices_coalesced, sorted_indices))
+        self.assertTrue(torch.equal(coalesce_permutation, permutation))
+
+    def test_batched_sort(self):
+        nr, nc = 4, 4
+        batch_size = 3
+        indices = torch.randperm(nr * nc, device=self.device)
+        indices = torch.stack([indices // nc, indices % nc])
+        sparse_indices = torch.cat([indices] * batch_size, dim=-1)
+        batch_indices = torch.arange(batch_size, device=self.device).repeat(16).unsqueeze(0)
+        batched_sparse_indices = torch.cat([batch_indices, sparse_indices])
+
+        values = torch.arange(nr * nc * batch_size, device=self.device)
+        sorted_indices_coalesced = torch.sparse_coo_tensor(batched_sparse_indices, values).coalesce().indices()
+        coalesce_permutation = torch.sparse_coo_tensor(batched_sparse_indices, values).coalesce().values()
+
+        sorted_indices, permutation = _sort_coo_indices(batched_sparse_indices)
+        self.assertTrue(torch.equal(sorted_indices_coalesced, sorted_indices))
+        self.assertTrue(torch.equal(coalesce_permutation, permutation))
+
+
+@parameterized_class(
+    (
+        "name",
+        "device",
+    ),
+    [
+        ("CPU", torch.device("cpu")),
+        (
+            "CUDA",
+            torch.device("cuda"),
+        ),
+    ],
+)
+class TestCOOtoCSR(unittest.TestCase):
+    def setUp(self) -> None:
+        if not torch.cuda.is_available() and self.device == torch.device("cuda"):
+            self.skipTest(f"Skipping {self.__class__.__name__} since CUDA is not available")
+
+    def batched_coo_to_csr(self, A_coo):
+        """
+        Converts batched sparse COO matrix A with shape [B, N, M]
+        to batched sparse CSR matrix with shape [B, N, M]
+
+        Inneficient implementation as the COO tensors only support int64 indices.
+        Meaning that int32 indices cannot be maintained if a CSR matrix is created via to_sparse_csr() from COO.
+        """
+        A_crow_indices_list = []
+        A_row_indices_list = []
+        A_values_list = []
+
+        size = A_coo.size()
+
+        for a_coo in A_coo:
+            a_csr = a_coo.detach().to_sparse_csr()  # detach to prevent grad on indices
+            A_crow_indices_list.append(a_csr.crow_indices())
+            A_row_indices_list.append(a_csr.col_indices())
+            A_values_list.append(a_csr.values())
+
+        A_crow_indices = torch.stack(A_crow_indices_list, dim=0)
+        A_row_indices = torch.stack(A_row_indices_list, dim=0)
+        A_values = torch.stack(A_values_list, dim=0)
+
+        return torch.sparse_csr_tensor(A_crow_indices, A_row_indices, A_values, size=size)
+
+    @parameterized.expand(
+        [
+            ("4x4_12n", torch.Size([4, 4]), 12),
+            ("8x16_32n", torch.Size([8, 16]), 32),
+        ]
+    )
+    def test_compress_row_indices(self, _, size, nnz):
+        A_coo = generate_random_sparse_coo_matrix(size, nnz, device=self.device)
+        A_csr = A_coo.to_sparse_csr()
+        A_csr_crow_indices = A_csr.crow_indices()
+        crow_indices = _compress_row_indices(A_coo.indices()[0], A_coo.size()[0])
+        self.assertTrue(torch.equal(A_csr_crow_indices, crow_indices))
+
+    @parameterized.expand(
+        [
+            ("4x4_12n", torch.Size([4, 4]), 12),
+            ("8x16_32n", torch.Size([8, 16]), 32),
+            ("4x4x4_12n", torch.Size([4, 4, 4]), 12),
+            ("6x8x14_32n", torch.Size([6, 8, 14]), 32),
+        ]
+    )
+    def test_coo_to_csr_indices(self, _, size, nnz):
+        A_coo = generate_random_sparse_coo_matrix(size, nnz, device=self.device)
+        A_csr = convert_coo_to_csr(A_coo)
+        if len(size) == 2:
+            A_csr_2 = A_coo.to_sparse_csr()
+        elif len(size) == 3:
+            A_csr_2 = self.batched_coo_to_csr(A_coo)
+        else:
+            raise ValueError(f"Size {size} not supported")
+
+        self.assertTrue(torch.equal(A_csr.crow_indices(), A_csr_2.crow_indices()))
+        self.assertTrue(torch.equal(A_csr.col_indices(), A_csr_2.col_indices()))
+
+    @parameterized.expand(
+        [
+            ("4x4_12n", torch.Size([4, 4]), 12),
+            ("8x16_32n", torch.Size([8, 16]), 32),
+            ("4x4x4_12n", torch.Size([4, 4, 4]), 12),
+            ("6x8x14_32n", torch.Size([6, 8, 14]), 32),
+        ]
+    )
+    def test_coo_to_csr_values(self, _, size, nnz):
+        A_coo = generate_random_sparse_coo_matrix(size, nnz, device=self.device)
+        A_csr = convert_coo_to_csr(A_coo)
+        if len(size) == 2:
+            A_csr_2 = A_coo.to_sparse_csr()
+        elif len(size) == 3:
+            A_csr_2 = self.batched_coo_to_csr(A_coo)
+        else:
+            raise ValueError(f"Size {size} not supported")
+
+        self.assertTrue(torch.equal(A_csr.values(), A_csr_2.values()))
+
+
+@parameterized_class(
+    (
+        "name",
+        "device",
+    ),
+    [
+        ("CPU", torch.device("cpu")),
+        (
+            "CUDA",
+            torch.device("cuda"),
+        ),
+    ],
+)
+class TestCSRtoCOO(unittest.TestCase):
+    def setUp(self) -> None:
+        if not torch.cuda.is_available() and self.device == torch.device("cuda"):
+            self.skipTest(f"Skipping {self.__class__.__name__} since CUDA is not available")
+
+    @parameterized.expand(
+        [
+            ("4x4_12n", torch.Size([4, 4]), 12),
+            ("8x16_32n", torch.Size([8, 16]), 32),
+        ]
+    )
+    def test_compress_row_indices(self, _, size, nnz):
+        A_csr = generate_random_sparse_csr_matrix(size, nnz, device=self.device)
+        A_coo = A_csr.to_sparse_coo()
+        A_coo_row_indices = A_coo.indices()[0]
+        row_indices = _demcompress_crow_indices(A_csr.crow_indices(), A_coo.size()[0])
+        self.assertTrue(torch.equal(A_coo_row_indices, row_indices))