[Sparse] Support sparse format conversion (#4929)

* [Sparse] Support sparse format conversion

* Update

* Update

dgl_sparse/include/sparse/sparse_format.h

@@ -18,42 +18,68 @@
 namespace dgl {
 namespace sparse {
 
-/** @brief SparseFormat enumeration */
+/** @brief SparseFormat enumeration. */
 enum SparseFormat { kCOO, kCSR, kCSC };
 
-/** @brief CSR sparse structure */
+/** @brief COO sparse structure. */
+struct COO {
+  /** @brief The shape of the matrix. */
+  int64_t num_rows = 0, num_cols = 0;
+  /** @brief COO format row indices array of the matrix. */
+  torch::Tensor row;
+  /** @brief COO format column indices array of the matrix. */
+  torch::Tensor col;
+  /** @brief Whether the row indices are sorted. */
+  bool row_sorted = false;
+  /** @brief Whether the column indices per row are sorted. */
+  bool col_sorted = false;
+};
+
+
+/** @brief CSR sparse structure. */
 struct CSR {
-  // CSR format index pointer array of the matrix
+  /** @brief The dense shape of the matrix. */
+  int64_t num_rows = 0, num_cols = 0;
+  /** @brief CSR format index pointer array of the matrix. */
   torch::Tensor indptr;
-  // CSR format index array of the matrix
+  /** @brief CSR format index array of the matrix. */
   torch::Tensor indices;
-  // The element order of the sparse format. In the SparseMatrix, we have data
-  // (value_) for each non-zero value. The order of non-zero values in (value_)
-  // may differ from the order of non-zero entries in CSR. So we store
-  // `value_indices` in CSR to indicate its relative non-zero value order to the
-  // SparseMatrix. With `value_indices`, we can retrieve the correct value for
-  // CSR, i.e., `value_[value_indices]`. If `value_indices` is not defined, this
-  // CSR follows the same non-zero value order as the SparseMatrix.
+  /** @brief Data index tensor. When it is null, assume it is from 0 to NNZ - 1.
+   */
   torch::optional<torch::Tensor> value_indices;
+  /** @brief Whether the column indices per row are sorted. */
+  bool sorted = false;
 };
 
-/** @brief COO sparse structure */
-struct COO {
-  // COO format row array of the matrix
-  torch::Tensor row;
-  // COO format column array of the matrix
-  torch::Tensor col;
-};
+/** @brief Convert an old DGL COO format to a COO in the sparse library. */
+std::shared_ptr<COO> COOFromOldDGLCOO(const aten::COOMatrix& dgl_coo);
 
-/**
- * @brief Convert a CSR format to COO format
- * @param num_rows Number of rows of the sparse format
- * @param num_cols Number of cols of the sparse format
- * @param csr CSR sparse format
- * @return COO sparse format
- */
-std::shared_ptr<COO> CSRToCOO(
-    int64_t num_rows, int64_t num_cols, const std::shared_ptr<CSR> csr);
+/** @brief Convert a COO in the sparse library to an old DGL COO matrix. */
+aten::COOMatrix COOToOldDGLCOO(const std::shared_ptr<COO>& coo);
+
+/** @brief Convert an old DGL CSR format to a CSR in the sparse library. */
+std::shared_ptr<CSR> CSRFromOldDGLCSR(const aten::CSRMatrix& dgl_csr);
+
+/** @brief Convert a CSR in the sparse library to an old DGL CSR matrix. */
+aten::CSRMatrix CSRToOldDGLCSR(const std::shared_ptr<CSR>& csr);
+
+/** @brief Convert a CSR format to COO format. */
+std::shared_ptr<COO> CSRToCOO(const std::shared_ptr<CSR>& csr);
+
+/** @brief Convert a CSC format to COO format. */
+std::shared_ptr<COO> CSCToCOO(const std::shared_ptr<CSR>& csc);
+
+/** @brief Convert a COO format to CSR format. */
+std::shared_ptr<CSR> COOToCSR(const std::shared_ptr<COO>& coo);
+
+/** @brief Convert a CSC format to CSR format. */
+std::shared_ptr<CSR> CSCToCSR(const std::shared_ptr<CSR>& csc);
+
+/** @brief Convert a COO format to CSC format. */
+std::shared_ptr<CSR> COOToCSC(const std::shared_ptr<COO>& coo);
+
+/** @brief Convert a CSR format to CSC format. */
+std::shared_ptr<CSR> CSRToCSC(const std::shared_ptr<CSR>& csr);
 
 }  // namespace sparse
 }  // namespace dgl

dgl_sparse/include/sparse/sparse_matrix.h

@@ -15,12 +15,14 @@
 #include <torch/script.h>
 
 #include <memory>
+#include <tuple>
+#include <utility>
 #include <vector>
 
 namespace dgl {
 namespace sparse {
 
-/** @brief SparseMatrix bound to Python  */
+/** @brief SparseMatrix bound to Python.  */
 class SparseMatrix : public torch::CustomClassHolder {
  public:
   /**
@@ -100,12 +102,14 @@ class SparseMatrix : public torch::CustomClassHolder {
   /** @brief Check whether this sparse matrix has CSC format. */
   inline bool HasCSC() const { return csc_ != nullptr; }
 
-  /** @return {row, col, value} tensors in the COO format. */
-  std::vector<torch::Tensor> COOTensors();
-  /** @return {row, col, value} tensors in the CSR format. */
-  std::vector<torch::Tensor> CSRTensors();
-  /** @return {row, col, value} tensors in the CSC format. */
-  std::vector<torch::Tensor> CSCTensors();
+  /** @return {row, col} tensors in the COO format. */
+  std::tuple<torch::Tensor, torch::Tensor> COOTensors();
+  /** @return {row, col, value_indices} tensors in the CSR format. */
+  std::tuple<torch::Tensor, torch::Tensor, torch::optional<torch::Tensor>>
+  CSRTensors();
+  /** @return {row, col, value_indices} tensors in the CSC format. */
+  std::tuple<torch::Tensor, torch::Tensor, torch::optional<torch::Tensor>>
+  CSCTensors();
 
   /**
    * @brief Set non-zero values of the sparse matrix

dgl_sparse/src/sparse_format.cc

@@ -14,21 +14,76 @@
 namespace dgl {
 namespace sparse {
 
-std::shared_ptr<COO> CSRToCOO(
-    int64_t num_rows, int64_t num_cols, const std::shared_ptr<CSR> csr) {
-  auto indptr = TorchTensorToDGLArray(csr->indptr);
-  auto indices = TorchTensorToDGLArray(csr->indices);
-  bool data_as_order = false;
-  runtime::NDArray data = aten::NullArray();
-  if (csr->value_indices.has_value()) {
-    data_as_order = true;
-    data = TorchTensorToDGLArray(csr->value_indices.value());
-  }
-  auto dgl_csr = aten::CSRMatrix(num_rows, num_cols, indptr, indices, data);
-  auto dgl_coo = aten::CSRToCOO(dgl_csr, data_as_order);
+std::shared_ptr<COO> COOFromOldDGLCOO(const aten::COOMatrix& dgl_coo) {
   auto row = DGLArrayToTorchTensor(dgl_coo.row);
   auto col = DGLArrayToTorchTensor(dgl_coo.col);
-  return std::make_shared<COO>(COO{row, col});
+  CHECK(aten::IsNullArray(dgl_coo.data));
+  return std::make_shared<COO>(
+      COO{dgl_coo.num_rows, dgl_coo.num_cols, row, col, dgl_coo.row_sorted,
+          dgl_coo.col_sorted});
+}
+
+aten::COOMatrix COOToOldDGLCOO(const std::shared_ptr<COO>& coo) {
+  auto row = TorchTensorToDGLArray(coo->row);
+  auto col = TorchTensorToDGLArray(coo->col);
+  return aten::COOMatrix(
+      coo->num_rows, coo->num_cols, row, col, aten::NullArray(),
+      coo->row_sorted, coo->col_sorted);
+}
+
+std::shared_ptr<CSR> CSRFromOldDGLCSR(const aten::CSRMatrix& dgl_csr) {
+  auto indptr = DGLArrayToTorchTensor(dgl_csr.indptr);
+  auto indices = DGLArrayToTorchTensor(dgl_csr.indices);
+  auto value_indices = DGLArrayToOptionalTorchTensor(dgl_csr.data);
+  return std::make_shared<CSR>(
+      CSR{dgl_csr.num_rows, dgl_csr.num_cols, indptr, indices, value_indices,
+          dgl_csr.sorted});
+}
+
+aten::CSRMatrix CSRToOldDGLCSR(const std::shared_ptr<CSR>& csr) {
+  auto indptr = TorchTensorToDGLArray(csr->indptr);
+  auto indices = TorchTensorToDGLArray(csr->indices);
+  auto data = OptionalTorchTensorToDGLArray(csr->value_indices);
+  return aten::CSRMatrix(
+      csr->num_rows, csr->num_cols, indptr, indices, data, csr->sorted);
+}
+
+std::shared_ptr<COO> CSRToCOO(const std::shared_ptr<CSR>& csr) {
+  auto dgl_csr = CSRToOldDGLCSR(csr);
+  auto dgl_coo = aten::CSRToCOO(dgl_csr, csr->value_indices.has_value());
+  return COOFromOldDGLCOO(dgl_coo);
+}
+
+std::shared_ptr<COO> CSCToCOO(const std::shared_ptr<CSR>& csc) {
+  auto dgl_csc = CSRToOldDGLCSR(csc);
+  auto dgl_coo = aten::CSRToCOO(dgl_csc, csc->value_indices.has_value());
+  dgl_coo = aten::COOTranspose(dgl_coo);
+  return COOFromOldDGLCOO(dgl_coo);
+}
+
+std::shared_ptr<CSR> COOToCSR(const std::shared_ptr<COO>& coo) {
+  auto dgl_coo = COOToOldDGLCOO(coo);
+  auto dgl_csr = aten::COOToCSR(dgl_coo);
+  return CSRFromOldDGLCSR(dgl_csr);
+}
+
+std::shared_ptr<CSR> CSCToCSR(const std::shared_ptr<CSR>& csc) {
+  auto dgl_csc = CSRToOldDGLCSR(csc);
+  auto dgl_csr = aten::CSRTranspose(dgl_csc);
+  return CSRFromOldDGLCSR(dgl_csr);
+}
+
+std::shared_ptr<CSR> COOToCSC(const std::shared_ptr<COO>& coo) {
+  auto dgl_coo = COOToOldDGLCOO(coo);
+  auto dgl_coo_transpose = aten::COOTranspose(dgl_coo);
+  auto dgl_csc = aten::COOToCSR(dgl_coo_transpose);
+  return CSRFromOldDGLCSR(dgl_csc);
+}
+
+std::shared_ptr<CSR> CSRToCSC(const std::shared_ptr<CSR>& csr) {
+  auto dgl_csr = CSRToOldDGLCSR(csr);
+  auto dgl_csc = aten::CSRTranspose(dgl_csr);
+  return CSRFromOldDGLCSR(dgl_csc);
 }
 
 }  // namespace sparse

dgl_sparse/src/sparse_matrix.cc

@@ -93,52 +93,65 @@ std::shared_ptr<CSR> SparseMatrix::CSCPtr() {
   return csc_;
 }
 
-std::vector<torch::Tensor> SparseMatrix::COOTensors() {
+std::tuple<torch::Tensor, torch::Tensor> SparseMatrix::COOTensors() {
   auto coo = COOPtr();
   auto val = value();
-  return {coo->row, coo->col, val};
+  return {coo->row, coo->col};
 }
 
-std::vector<torch::Tensor> SparseMatrix::CSRTensors() {
+std::tuple<torch::Tensor, torch::Tensor, torch::optional<torch::Tensor>>
+SparseMatrix::CSRTensors() {
   auto csr = CSRPtr();
   auto val = value();
-  if (csr->value_indices.has_value()) {
-    val = val[csr->value_indices.value()];
-  }
-  return {csr->indptr, csr->indices, val};
+  return {csr->indptr, csr->indices, csr->value_indices};
 }
 
-std::vector<torch::Tensor> SparseMatrix::CSCTensors() {
+std::tuple<torch::Tensor, torch::Tensor, torch::optional<torch::Tensor>>
+SparseMatrix::CSCTensors() {
   auto csc = CSCPtr();
-  auto val = value();
-  if (csc->value_indices.has_value()) {
-    val = val[csc->value_indices.value()];
-  }
-  return {csc->indptr, csc->indices, val};
+  return {csc->indptr, csc->indices, csc->value_indices};
 }
 
 void SparseMatrix::SetValue(torch::Tensor value) { value_ = value; }
 
 void SparseMatrix::_CreateCOO() {
-  if (HasCOO()) {
-    return;
-  }
+  if (HasCOO()) return;
   if (HasCSR()) {
-    coo_ = CSRToCOO(shape_[0], shape_[1], csr_);
+    coo_ = CSRToCOO(csr_);
   } else if (HasCSC()) {
-    // TODO(zhenkun)
+    coo_ = CSCToCOO(csc_);
   } else {
     LOG(FATAL) << "SparseMatrix does not have any sparse format";
   }
 }
 
-void SparseMatrix::_CreateCSR() {}
-void SparseMatrix::_CreateCSC() {}
+void SparseMatrix::_CreateCSR() {
+  if (HasCSR()) return;
+  if (HasCOO()) {
+    csr_ = COOToCSR(coo_);
+  } else if (HasCSC()) {
+    csr_ = CSCToCSR(csc_);
+  } else {
+    LOG(FATAL) << "SparseMatrix does not have any sparse format";
+  }
+}
+
+void SparseMatrix::_CreateCSC() {
+  if (HasCSC()) return;
+  if (HasCOO()) {
+    csc_ = COOToCSC(coo_);
+  } else if (HasCSR()) {
+    csc_ = CSRToCSC(csr_);
+  } else {
+    LOG(FATAL) << "SparseMatrix does not have any sparse format";
+  }
+}
 
 c10::intrusive_ptr<SparseMatrix> CreateFromCOO(
     torch::Tensor row, torch::Tensor col, torch::Tensor value,
     const std::vector<int64_t>& shape) {
-  auto coo = std::make_shared<COO>(COO{row, col});
+  auto coo =
+      std::make_shared<COO>(COO{shape[0], shape[1], row, col, false, false});
   return SparseMatrix::FromCOO(coo, value, shape);
 }
 
@@ -146,7 +159,8 @@ c10::intrusive_ptr<SparseMatrix> CreateFromCSR(
     torch::Tensor indptr, torch::Tensor indices, torch::Tensor value,
     const std::vector<int64_t>& shape) {
   auto csr = std::make_shared<CSR>(
-      CSR{indptr, indices, torch::optional<torch::Tensor>()});
+      CSR{shape[0], shape[1], indptr, indices, torch::optional<torch::Tensor>(),
+          false});
   return SparseMatrix::FromCSR(csr, value, shape);
 }
 
@@ -154,7 +168,8 @@ c10::intrusive_ptr<SparseMatrix> CreateFromCSC(
     torch::Tensor indptr, torch::Tensor indices, torch::Tensor value,
     const std::vector<int64_t>& shape) {
   auto csc = std::make_shared<CSR>(
-      CSR{indptr, indices, torch::optional<torch::Tensor>()});
+      CSR{shape[1], shape[0], indptr, indices, torch::optional<torch::Tensor>(),
+          false});
   return SparseMatrix::FromCSC(csc, value, shape);
 }
 

dgl_sparse/src/utils.h

@@ -61,6 +61,24 @@ inline static torch::Tensor DGLArrayToTorchTensor(runtime::NDArray array) {
   return at::fromDLPack(runtime::DLPackConvert::ToDLPack(array));
 }
 
+/** @brief Convert an optional Torch tensor to a DGL array. */
+inline static runtime::NDArray OptionalTorchTensorToDGLArray(
+    torch::optional<torch::Tensor> tensor) {
+  if (!tensor.has_value()) {
+    return aten::NullArray();
+  }
+  return TorchTensorToDGLArray(tensor.value());
+}
+
+/** @brief Convert a DGL array to an optional Torch tensor. */
+inline static torch::optional<torch::Tensor> DGLArrayToOptionalTorchTensor(
+    runtime::NDArray array) {
+  if (aten::IsNullArray(array)) {
+    return torch::optional<torch::Tensor>();
+  }
+  return torch::make_optional<torch::Tensor>(DGLArrayToTorchTensor(array));
+}
+
 }  // namespace sparse
 }  // namespace dgl
 

python/dgl/mock_sparse2/sparse_matrix.py

@@ -109,8 +109,8 @@ class SparseMatrix:
 
         Returns
         -------
-        Tuple[torch.Tensor, torch.Tensor, torch.Tensor]
-            A tuple of tensors containing row, column coordinates and values.
+        Tuple[torch.Tensor, torch.Tensor]
+            A tuple of tensors containing row and column coordinates.
         """
         return self.c_sparse_matrix.coo()
 
@@ -120,7 +120,7 @@ class SparseMatrix:
         Returns
         -------
         Tuple[torch.Tensor, torch.Tensor, torch.Tensor]
-            A tuple of tensors containing row, column coordinates and values.
+            A tuple of tensors containing row, column coordinates and value indices.
         """
         return self.c_sparse_matrix.csr()
 
@@ -130,7 +130,7 @@ class SparseMatrix:
         Returns
         -------
         Tuple[torch.Tensor, torch.Tensor, torch.Tensor]
-            A tuple of tensors containing row, column coordinates and values.
+            A tuple of tensors containing row, column coordinates and value indices.
         """
         return self.c_sparse_matrix.csc()
 
@@ -142,7 +142,8 @@ class SparseMatrix:
         torch.Tensor
             Dense representation of the sparse matrix.
         """
-        row, col, val = self.coo()
+        row, col = self.coo()
+        val = self.val
         shape = self.shape + val.shape[1:]
         mat = torch.zeros(shape, device=self.device)
         mat[row, col] = val
@@ -186,16 +187,16 @@ def create_from_coo(
     >>> A = create_from_coo(dst, src)
     >>> print(A)
     SparseMatrix(indices=tensor([[1, 1, 2],
-                                 [2, 4, 3]]),
-                 values=tensor([1., 1., 1.]),
-                 shape=(3, 5), nnz=3)
+            [2, 4, 3]]),
+    values=tensor([1., 1., 1.]),
+    shape=(3, 5), nnz=3)
     >>> # Specify shape
     >>> A = create_from_coo(dst, src, shape=(5, 5))
     >>> print(A)
     SparseMatrix(indices=tensor([[1, 1, 2],
-                                 [2, 4, 3]]),
-                 values=tensor([1., 1., 1.]),
-                 shape=(5, 5), nnz=3)
+            [2, 4, 3]]),
+    values=tensor([1., 1., 1.]),
+    shape=(5, 5), nnz=3)
 
     Case2: Sparse matrix with scalar/vector values. Following example is with
     vector data.
@@ -203,11 +204,11 @@ def create_from_coo(
     >>> val = torch.tensor([[1., 1.], [2., 2.], [3., 3.]])
     >>> A = create_from_coo(dst, src, val)
     SparseMatrix(indices=tensor([[1, 1, 2],
-                                 [2, 4, 3]]),
-                 values=tensor([[1, 1],
-                                [2, 2],
-                                [3, 3]]),
-                 shape=(3, 5), nnz=3)
+            [2, 4, 3]]),
+    values=tensor([[1., 1.],
+            [2., 2.],
+            [3., 3.]]),
+    shape=(3, 5), nnz=3)
     """
     if shape is None:
         shape = (torch.max(row).item() + 1, torch.max(col).item() + 1)
@@ -219,8 +220,6 @@ def create_from_coo(
     )
 
 
-# FIXME: The docstring cannot print A because we cannot print
-# the indices of CSR/CSC
 def create_from_csr(
     indptr: torch.Tensor,
     indices: torch.Tensor,
@@ -269,16 +268,16 @@ def create_from_csr(
     >>> A = create_from_csr(indptr, indices)
     >>> print(A)
     SparseMatrix(indices=tensor([[0, 1, 2, 2, 2],
-                                 [1, 2, 0, 1, 2]]),
-                 values=tensor([1., 1., 1., 1., 1.]),
-                 shape=(3, 3), nnz=5)
+            [1, 2, 0, 1, 2]]),
+    values=tensor([1., 1., 1., 1., 1.]),
+    shape=(3, 3), nnz=5)
     >>> # Specify shape
-    >>> A = create_from_csr(indptr, indices, shape=(5, 3))
+    >>> A = create_from_csr(indptr, indices, shape=(3, 5))
     >>> print(A)
     SparseMatrix(indices=tensor([[0, 1, 2, 2, 2],
             [1, 2, 0, 1, 2]]),
     values=tensor([1., 1., 1., 1., 1.]),
-    shape=(5, 3), nnz=5)
+    shape=(3, 5), nnz=5)
 
     Case2: Sparse matrix with scalar/vector values. Following example is with
     vector data.
@@ -305,8 +304,6 @@ def create_from_csr(
     )
 
 
-# FIXME: The docstring cannot print A because we cannot print
-# the indices of CSR/CSC
 def create_from_csc(
     indptr: torch.Tensor,
     indices: torch.Tensor,
@@ -354,15 +351,15 @@ def create_from_csc(
     >>> indices = torch.tensor([2, 0, 2, 1, 2])
     >>> A = create_from_csc(indptr, indices)
     >>> print(A)
-    SparseMatrix(indices=tensor([[0, 1, 2, 2, 2],
-                                 [1, 2, 0, 1, 2]]),
-                 values=tensor([1., 1., 1., 1., 1.]),
-                 shape=(3, 3), nnz=5)
+    SparseMatrix(indices=tensor([[2, 0, 2, 1, 2],
+            [0, 1, 1, 2, 2]]),
+    values=tensor([1., 1., 1., 1., 1.]),
+    shape=(3, 3), nnz=5)
     >>> # Specify shape
     >>> A = create_from_csc(indptr, indices, shape=(5, 3))
     >>> print(A)
-    SparseMatrix(indices=tensor([[0, 1, 2, 2, 2],
-            [1, 2, 0, 1, 2]]),
+    SparseMatrix(indices=tensor([[2, 0, 2, 1, 2],
+            [0, 1, 1, 2, 2]]),
     values=tensor([1., 1., 1., 1., 1.]),
     shape=(5, 3), nnz=5)
 
@@ -372,12 +369,12 @@ def create_from_csc(
     >>> val = torch.tensor([[1, 1], [2, 2], [3, 3], [4, 4], [5, 5]])
     >>> A = create_from_csc(indptr, indices, val)
     >>> print(A)
-    SparseMatrix(indices=tensor([[0, 1, 2, 2, 2],
-            [1, 2, 0, 1, 2]]),
-    values=tensor([[2, 2],
-            [4, 4],
-            [1, 1],
+    SparseMatrix(indices=tensor([[2, 0, 2, 1, 2],
+            [0, 1, 1, 2, 2]]),
+    values=tensor([[1, 1],
+            [2, 2],
             [3, 3],
+            [4, 4],
             [5, 5]]),
     shape=(3, 3), nnz=5)
     """

tests/pytorch/mock_sparse2/test_diag.py

@@ -44,7 +44,8 @@ def test_diag(val_shape, mat_shape):
     assert sp_mat.device == mat.device
     # row, col, val
     edge_index = torch.arange(len(val)).to(mat.device)
-    row, col, val = sp_mat.coo()
+    row, col = sp_mat.coo()
+    val = sp_mat.val
     assert torch.allclose(row, edge_index)
     assert torch.allclose(col, edge_index)
     assert torch.allclose(val, val)

tests/pytorch/mock_sparse2/test_elementwise_op_sp.py

@@ -12,7 +12,8 @@ if not sys.platform.startswith("linux"):
     pytest.skip("skipping tests on win", allow_module_level=True)
 
 def all_close_sparse(A, row, col, val, shape):
-    rowA, colA, valA = A.coo()
+    rowA, colA = A.coo()
+    valA = A.val
     assert torch.allclose(rowA, row)
     assert torch.allclose(colA, col)
     assert torch.allclose(valA, val)

tests/pytorch/mock_sparse2/test_sparse_matrix.py

@@ -10,17 +10,12 @@ from dgl.mock_sparse2 import create_from_coo, create_from_csr, create_from_csc
 if not sys.platform.startswith("linux"):
     pytest.skip("skipping tests on win", allow_module_level=True)
 
+
 @pytest.mark.parametrize("dense_dim", [None, 4])
-@pytest.mark.parametrize("row", [[0, 0, 1, 2], (0, 1, 2, 4)])
+@pytest.mark.parametrize("row", [(0, 0, 1, 2), (0, 1, 2, 4)])
 @pytest.mark.parametrize("col", [(0, 1, 2, 2), (1, 3, 3, 4)])
-@pytest.mark.parametrize("shape", [None, (3, 5), (5, 3)])
+@pytest.mark.parametrize("shape", [None, (5, 5), (5, 6)])
 def test_create_from_coo(dense_dim, row, col, shape):
-    # Skip invalid matrices
-    if shape is not None and (
-        max(row) >= shape[0] or max(col) >= shape[1]
-    ):
-        return
-
     val_shape = (len(row),)
     if dense_dim is not None:
         val_shape += (dense_dim,)
@@ -33,7 +28,8 @@ def test_create_from_coo(dense_dim, row, col, shape):
     if shape is None:
         shape = (torch.max(row).item() + 1, torch.max(col).item() + 1)
 
-    mat_row, mat_col, mat_val = mat.coo()
+    mat_row, mat_col = mat.coo()
+    mat_val = mat.val
 
     assert mat.shape == shape
     assert mat.nnz == row.numel()
@@ -44,7 +40,7 @@ def test_create_from_coo(dense_dim, row, col, shape):
 
 
 @pytest.mark.parametrize("dense_dim", [None, 4])
-@pytest.mark.parametrize("indptr", [[0, 0, 1, 4], (0, 1, 2, 4)])
+@pytest.mark.parametrize("indptr", [(0, 0, 1, 4), (0, 1, 2, 4)])
 @pytest.mark.parametrize("indices", [(0, 1, 2, 3), (1, 2, 3, 4)])
 @pytest.mark.parametrize("shape", [None, (3, 5)])
 def test_create_from_csr(dense_dim, indptr, indices, shape):
@@ -64,13 +60,15 @@ def test_create_from_csr(dense_dim, indptr, indices, shape):
     assert mat.shape == shape
     assert mat.nnz == indices.numel()
     assert mat.dtype == val.dtype
-    mat_indptr, mat_indices, mat_val = mat.csr()
+    mat_indptr, mat_indices, value_indices = mat.csr()
+    mat_val = mat.val if value_indices is None else mat.val[value_indices]
     assert torch.allclose(mat_indptr, indptr)
     assert torch.allclose(mat_indices, indices)
     assert torch.allclose(mat_val, val)
 
+
 @pytest.mark.parametrize("dense_dim", [None, 4])
-@pytest.mark.parametrize("indptr", [[0, 0, 1, 4], (0, 1, 2, 4)])
+@pytest.mark.parametrize("indptr", [(0, 0, 1, 4), (0, 1, 2, 4)])
 @pytest.mark.parametrize("indices", [(0, 1, 2, 3), (1, 2, 3, 4)])
 @pytest.mark.parametrize("shape", [None, (5, 3)])
 def test_create_from_csc(dense_dim, indptr, indices, shape):
@@ -90,11 +88,13 @@ def test_create_from_csc(dense_dim, indptr, indices, shape):
     assert mat.shape == shape
     assert mat.nnz == indices.numel()
     assert mat.dtype == val.dtype
-    mat_indptr, mat_indices, mat_val = mat.csc()
+    mat_indptr, mat_indices, value_indices = mat.csc()
+    mat_val = mat.val if value_indices is None else mat.val[value_indices]
     assert torch.allclose(mat_indptr, indptr)
     assert torch.allclose(mat_indices, indices)
     assert torch.allclose(mat_val, val)
 
+
 @pytest.mark.parametrize("val_shape", [(3), (3, 2)])
 def test_dense(val_shape):
     ctx = F.ctx()
@@ -110,6 +110,7 @@ def test_dense(val_shape):
     mat[row, col] = val
     assert torch.allclose(A_dense, mat)
 
+
 def test_set_val():
     ctx = F.ctx()
 
@@ -123,8 +124,9 @@ def test_set_val():
     A.val = new_val
     assert torch.allclose(new_val, A.val)
 
+
 @pytest.mark.parametrize("dense_dim", [None, 4])
-@pytest.mark.parametrize("indptr", [[0, 0, 1, 4], (0, 1, 2, 4)])
+@pytest.mark.parametrize("indptr", [(0, 0, 1, 4), (0, 1, 2, 4)])
 @pytest.mark.parametrize("indices", [(0, 1, 2, 3), (1, 4, 3, 2)])
 @pytest.mark.parametrize("shape", [None, (3, 5)])
 def test_csr_to_coo(dense_dim, indptr, indices, shape):
@@ -140,9 +142,14 @@ def test_csr_to_coo(dense_dim, indptr, indices, shape):
     if shape is None:
         shape = (indptr.numel() - 1, torch.max(indices).item() + 1)
 
-    row = torch.arange(0, indptr.shape[0] - 1).to(ctx).repeat_interleave(torch.diff(indptr))
+    row = (
+        torch.arange(0, indptr.shape[0] - 1)
+        .to(ctx)
+        .repeat_interleave(torch.diff(indptr))
+    )
     col = indices
-    mat_row, mat_col, mat_val = mat.coo()
+    mat_row, mat_col = mat.coo()
+    mat_val = mat.val
 
     assert mat.shape == shape
     assert mat.nnz == row.numel()
@@ -152,3 +159,194 @@ def test_csr_to_coo(dense_dim, indptr, indices, shape):
     assert torch.allclose(mat_row, row)
     assert torch.allclose(mat_col, col)
 
+
+@pytest.mark.parametrize("dense_dim", [None, 4])
+@pytest.mark.parametrize("indptr", [(0, 0, 1, 4), (0, 1, 2, 4)])
+@pytest.mark.parametrize("indices", [(0, 1, 2, 3), (1, 4, 3, 2)])
+@pytest.mark.parametrize("shape", [None, (5, 3)])
+def test_csc_to_coo(dense_dim, indptr, indices, shape):
+    ctx = F.ctx()
+    val_shape = (len(indices),)
+    if dense_dim is not None:
+        val_shape += (dense_dim,)
+    val = torch.randn(val_shape).to(ctx)
+    indptr = torch.tensor(indptr).to(ctx)
+    indices = torch.tensor(indices).to(ctx)
+    mat = create_from_csc(indptr, indices, val, shape)
+
+    if shape is None:
+        shape = (torch.max(indices).item() + 1, indptr.numel() - 1)
+
+    col = (
+        torch.arange(0, indptr.shape[0] - 1)
+        .to(ctx)
+        .repeat_interleave(torch.diff(indptr))
+    )
+    row = indices
+    mat_row, mat_col = mat.coo()
+    mat_val = mat.val
+
+    assert mat.shape == shape
+    assert mat.nnz == row.numel()
+    assert mat.device == row.device
+    assert mat.dtype == val.dtype
+    assert torch.allclose(mat_val, val)
+    assert torch.allclose(mat_row, row)
+    assert torch.allclose(mat_col, col)
+
+
+def _scatter_add(a, index, v=1):
+    index = index.tolist()
+    for i in index:
+        a[i] += v
+    return a
+
+
+@pytest.mark.parametrize("dense_dim", [None, 4])
+@pytest.mark.parametrize("row", [(0, 0, 1, 2), (0, 1, 2, 4)])
+@pytest.mark.parametrize("col", [(0, 1, 2, 2), (1, 3, 3, 4)])
+@pytest.mark.parametrize("shape", [None, (5, 5), (5, 6)])
+def test_coo_to_csr(dense_dim, row, col, shape):
+    val_shape = (len(row),)
+    if dense_dim is not None:
+        val_shape += (dense_dim,)
+    ctx = F.ctx()
+    val = torch.randn(val_shape).to(ctx)
+    row = torch.tensor(row).to(ctx)
+    col = torch.tensor(col).to(ctx)
+    mat = create_from_coo(row, col, val, shape)
+
+    if shape is None:
+        shape = (torch.max(row).item() + 1, torch.max(col).item() + 1)
+
+    mat_indptr, mat_indices, value_indices = mat.csr()
+    mat_val = mat.val if value_indices is None else mat.val[value_indices]
+    indptr = torch.zeros(shape[0] + 1).to(ctx)
+    indptr = _scatter_add(indptr, row + 1)
+    indptr = torch.cumsum(indptr, 0).long()
+    indices = col
+
+    assert mat.shape == shape
+    assert mat.nnz == row.numel()
+    assert mat.dtype == val.dtype
+    assert torch.allclose(mat_val, val)
+    assert torch.allclose(mat_indptr, indptr)
+    assert torch.allclose(mat_indices, indices)
+
+
+@pytest.mark.parametrize("dense_dim", [None, 4])
+@pytest.mark.parametrize("indptr", [(0, 0, 1, 4), (0, 1, 2, 4)])
+@pytest.mark.parametrize("indices", [(0, 1, 2, 3), (1, 4, 3, 2)])
+@pytest.mark.parametrize("shape", [None, (5, 3)])
+def test_csc_to_csr(dense_dim, indptr, indices, shape):
+    ctx = F.ctx()
+    val_shape = (len(indices),)
+    if dense_dim is not None:
+        val_shape += (dense_dim,)
+    val = torch.randn(val_shape).to(ctx)
+    indptr = torch.tensor(indptr).to(ctx)
+    indices = torch.tensor(indices).to(ctx)
+    mat = create_from_csc(indptr, indices, val, shape)
+    mat_indptr, mat_indices, value_indices = mat.csr()
+    mat_val = mat.val if value_indices is None else mat.val[value_indices]
+
+    if shape is None:
+        shape = (torch.max(indices).item() + 1, indptr.numel() - 1)
+
+    col = (
+        torch.arange(0, indptr.shape[0] - 1)
+        .to(ctx)
+        .repeat_interleave(torch.diff(indptr))
+    )
+    row = indices
+    row, sort_index = row.sort(stable=True)
+    col = col[sort_index]
+    val = val[sort_index]
+    indptr = torch.zeros(shape[0] + 1).to(ctx)
+    indptr = _scatter_add(indptr, row + 1)
+    indptr = torch.cumsum(indptr, 0).long()
+    indices = col
+
+    assert mat.shape == shape
+    assert mat.nnz == row.numel()
+    assert mat.device == row.device
+    assert mat.dtype == val.dtype
+    assert torch.allclose(mat_val, val)
+    assert torch.allclose(mat_indptr, indptr)
+    assert torch.allclose(mat_indices, indices)
+
+
+@pytest.mark.parametrize("dense_dim", [None, 4])
+@pytest.mark.parametrize("row", [(0, 0, 1, 2), (0, 1, 2, 4)])
+@pytest.mark.parametrize("col", [(0, 1, 2, 2), (1, 3, 3, 4)])
+@pytest.mark.parametrize("shape", [None, (5, 5), (5, 6)])
+def test_coo_to_csc(dense_dim, row, col, shape):
+
+    val_shape = (len(row),)
+    if dense_dim is not None:
+        val_shape += (dense_dim,)
+    ctx = F.ctx()
+    val = torch.randn(val_shape).to(ctx)
+    row = torch.tensor(row).to(ctx)
+    col = torch.tensor(col).to(ctx)
+    mat = create_from_coo(row, col, val, shape)
+
+    if shape is None:
+        shape = (torch.max(row).item() + 1, torch.max(col).item() + 1)
+
+    mat_indptr, mat_indices, value_indices = mat.csc()
+    mat_val = mat.val if value_indices is None else mat.val[value_indices]
+    indptr = torch.zeros(shape[1] + 1).to(ctx)
+    _scatter_add(indptr, col + 1)
+    indptr = torch.cumsum(indptr, 0).long()
+    indices = row
+
+    assert mat.shape == shape
+    assert mat.nnz == row.numel()
+    assert mat.dtype == val.dtype
+    assert torch.allclose(mat_val, val)
+    assert torch.allclose(mat_indptr, indptr)
+    assert torch.allclose(mat_indices, indices)
+
+
+@pytest.mark.parametrize("dense_dim", [None, 4])
+@pytest.mark.parametrize("indptr", [(0, 0, 1, 4), (0, 1, 2, 4)])
+@pytest.mark.parametrize("indices", [(0, 1, 2, 3), (1, 2, 3, 4)])
+@pytest.mark.parametrize("shape", [None, (3, 5)])
+def test_csr_to_csc(dense_dim, indptr, indices, shape):
+    val_shape = (len(indices),)
+    if dense_dim is not None:
+        val_shape += (dense_dim,)
+    ctx = F.ctx()
+    val = torch.randn(val_shape).to(ctx)
+    indptr = torch.tensor(indptr).to(ctx)
+    indices = torch.tensor(indices).to(ctx)
+    mat = create_from_csr(indptr, indices, val, shape)
+    mat_indptr, mat_indices, value_indices = mat.csc()
+    mat_val = mat.val if value_indices is None else mat.val[value_indices]
+
+    if shape is None:
+        shape = (indptr.numel() - 1, torch.max(indices).item() + 1)
+
+    row = (
+        torch.arange(0, indptr.shape[0] - 1)
+        .to(ctx)
+        .repeat_interleave(torch.diff(indptr))
+    )
+
+    col = indices
+    col, sort_index = col.sort(stable=True)
+    row = row[sort_index]
+    val = val[sort_index]
+    indptr = torch.zeros(shape[1] + 1).to(ctx)
+    indptr = _scatter_add(indptr, col + 1)
+    indptr = torch.cumsum(indptr, 0).long()
+    indices = row
+
+    assert mat.shape == shape
+    assert mat.nnz == row.numel()
+    assert mat.device == row.device
+    assert mat.dtype == val.dtype
+    assert torch.allclose(mat_val, val)
+    assert torch.allclose(mat_indptr, indptr)
+    assert torch.allclose(mat_indices, indices)