Repository: dmlc/dgl
Pull request: Move mock version of dgl_sparse library to DGL main repo
## Description
Resolves issue https://github.com/dmlc/dgl/issues/4489
This PR is  a duplicate of [PR #4500](https://github.com/dmlc/dgl/pull/4500). The old PR was created on the base repo `dmlc/dgl` (not on a forked repo) by mistake. 

## Checklist
Please feel free to remove inapplicable items for your PR.
- [ ] The PR title starts with [$CATEGORY] (such as [NN], [Model], [Doc], [Feature]])
- [ ] Changes are complete (i.e. I finished coding on this PR)
- [ ] All changes have test coverage
- [ ] Code is well-documented
- [ ] To the best of my knowledge, examples are either not affected by this change,
      or have been fixed to be compatible with this change
- [ ] Related issue is referred in this PR
- [ ] If the PR is for a new model/paper, I've updated the example index [here](../examples/README.md).

## Changes
<!-- You could use following template
- [ ] Feature1, tests, (and when applicable, API doc)
- [ ] Feature2, tests, (and when applicable, API doc)
-->


Problem statement:
Sparse Ops: Move code from temporary repo to DGL main repo
Move the codes developed in the temporary repo to DGL main repo (`python/dgl/mock_sparse`) and setup:
- Unit test & CI
- Pass lint check
- Doc

The subpackage shall NOT be imported by the main package by default to avoid breaking existing logics.

Hints from issue discussion:

Commit message: rm docstring for attributes for api doc

Patch to review:
diff --git a/docs/source/api/python/dgl.sparse_v0.rst b/docs/source/api/python/dgl.sparse_v0.rst
new file mode 100644
index 000000000000..149cc474e736
--- /dev/null
+++ b/docs/source/api/python/dgl.sparse_v0.rst
@@ -0,0 +1,14 @@
+.. _apibackend:
+
+dgl.sparse
+=================================
+
+`dgl_sparse` is a library for sparse operators that are commonly used in GNN models.
+
+
+Sparse matrix class
+-------------------------
+.. currentmodule:: dgl.mock_sparse
+
+.. autoclass:: SparseMatrix
+	:members:
\ No newline at end of file
diff --git a/python/dgl/mock_sparse/__init__.py b/python/dgl/mock_sparse/__init__.py
new file mode 100644
index 000000000000..0d5521c11b32
--- /dev/null
+++ b/python/dgl/mock_sparse/__init__.py
@@ -0,0 +1,2 @@
+"""dgl sparse class."""
+from .sp_matrix import SparseMatrix
diff --git a/python/dgl/mock_sparse/elementwise_op_sp.py b/python/dgl/mock_sparse/elementwise_op_sp.py
new file mode 100644
index 000000000000..ff6f9f9ad775
--- /dev/null
+++ b/python/dgl/mock_sparse/elementwise_op_sp.py
@@ -0,0 +1,219 @@
+"""dgl elementwise operators for sparse matrix module."""
+import torch
+from .sp_matrix import SparseMatrix
+
+def check_sparsity(A, B):
+    """Checks if two sparse matrices have the same sparsity structure."""
+    if not torch.equal(A.indices("COO"), B.indices("COO")):
+        raise ValueError('The two input matrices have different sparsity pattern')
+
+def add(A, B):
+    """Elementwise addition.
+
+    Parameters
+    ----------
+    A : SparseMatrix
+        Sparse matrix
+    B : SparseMatrix
+        Sparse matrix
+
+    Returns
+    -------
+    SparseMatrix
+        Sparse matrix
+
+    Examples
+    --------
+    >>> row = torch.tensor([1, 1, 2])
+    >>> col = torch.tensor([2, 4, 3])
+    >>> w1 = torch.tensor([3 ,4, 5])
+    >>> w2 = torch.tensor([-1, -3, -3])
+    >>> A = SparseMatrix(row, col, w1)
+    >>> C = A + A
+    >>> C.val
+    tensor([ 6,  8, 10])
+    >>> C = A(w2) + A(w1)
+    tensor([2, 1, 2])
+    """
+    if isinstance(A, SparseMatrix) and isinstance(B, SparseMatrix):
+        assert A.shape == B.shape, 'The shape of sparse matrix A {} and' \
+        'B {} must match'.format(A.shape, B.shape)
+        check_sparsity(A, B)
+        return SparseMatrix(A.row, A.col, A.val + B.val)
+    raise RuntimeError('Elementwise add between sparse and dense matrix is not supported.')
+
+def sub(A, B):
+    """Elementwise subtraction.
+
+    Parameters
+    ----------
+    A : SparseMatrix
+        Sparse matrix
+    B : SparseMatrix
+        Sparse matrix
+
+    Returns
+    -------
+    SparseMatrix
+        Sparse matrix
+
+    Examples
+    --------
+    ...
+    >>> C = A - A
+    >>> C.val
+    tensor([0, 0, 0])
+    >>> C = A(w2) - A(w1)
+    tensor([-4, -7, -8])
+    """
+    if isinstance(A, SparseMatrix) and isinstance(B, SparseMatrix):
+        assert A.shape == B.shape, 'The shape of sparse matrix A {} and' \
+        'B {} must match'.format(A.shape, B.shape)
+        check_sparsity(A, B)
+        return SparseMatrix(A.row, A.col, A.val - B.val)
+    raise RuntimeError('Elementwise sub between sparse and dense matrix is not supported.')
+
+def rsub(A, B):
+    """Elementwise subtraction.
+
+    Parameters
+    ----------
+    A : scalar
+        scalar value
+    B : SparseMatrix
+        Sparse matrix
+    """
+    raise RuntimeError('Elementwise sub between sparse and dense matrix is not supported.')
+
+def mul(A, B):
+    """Elementwise multiplication.
+
+    Parameters
+    ----------
+    A : SparseMatrix or scalar
+        Sparse matrix or scalar value
+    B : SparseMatrix or scalar
+        Sparse matrix or scalar value.
+
+    Returns
+    -------
+    SparseMatrix
+        Sparse matrix
+
+    Examples
+    --------
+    ...
+    >>> C = A * A
+    >>> C.val
+    tensor([ 9, 16, 25])
+    >>> C = A(w2) * A(w1)
+    tensor([ -3, -12, -15])
+    >>> v_scalar = 2.5
+    >>> C = A * v_scalar
+    >>> C.val
+    tensor([ 7.5000, 10.0000, 12.5000])
+    >>> C = v_scalar * A
+    >>> C.val
+    tensor([ 7.5000, 10.0000, 12.5000])
+    """
+    if isinstance(A, SparseMatrix) and isinstance(B, SparseMatrix):
+        check_sparsity(A, B)
+    a_val = A.val if isinstance(A, SparseMatrix) else A
+    b_val = B.val if isinstance(B, SparseMatrix) else B
+    return SparseMatrix(A.row, A.col, a_val * b_val)
+
+def div(A, B):
+    """Elementwise division.
+
+    Parameters
+    ----------
+    A : SparseMatrix
+        Sparse matrix
+    B : SparseMatrix or scalar
+        Sparse matrix or scalar value.
+
+    Returns
+    -------
+    SparseMatrix
+        Sparse matrix
+
+    Examples
+    --------
+    ...
+    >>> C = A / A
+    >>> C.val
+    tensor([1., 1., 1.])
+    >>> C = A(w2) / A(w1)
+    tensor([-0.3333, -0.7500, -0.6000])
+    >>> v_scalar = 2.5
+    >>> C = A / v_scalar
+    >>> C.val
+    tensor([1.2000, 1.6000, 2.0000])
+    """
+    if isinstance(A, SparseMatrix) and isinstance(B, SparseMatrix):
+        check_sparsity(A, B)
+    a_val = A.val if isinstance(A, SparseMatrix) else A
+    b_val = B.val if isinstance(B, SparseMatrix) else B
+    return SparseMatrix(A.row, A.col, a_val / b_val)
+
+def rdiv(A, B):
+    """Elementwise division.
+
+    Parameters
+    ----------
+    A : scalar
+        scalar value
+    B : SparseMatrix
+        Sparse matrix
+    """
+    raise RuntimeError('Elementwise div between sparse and dense matrix is not supported.')
+
+def power(A, B):
+    """Elementwise power operation.
+
+    Parameters
+    ----------
+    A : SparseMatrix
+        Sparse matrix
+    B : scalar
+        scalar value.
+
+    Returns
+    -------
+    SparseMatrix
+        Sparse matrix
+
+    Examples
+    --------
+    ...
+    >>> C = pow(A, v_scalar)
+    >>> C.val
+    tensor([15.5885, 32.0000, 55.9017])
+    """
+    if isinstance(B, SparseMatrix):
+        raise RuntimeError('power operation between two sparse matrices is not supported')
+    return SparseMatrix(A.row, A.col, torch.pow(A.val, B))
+
+def rpower(A, B):
+    """Elementwise subtraction.
+
+    Parameters
+    ----------
+    A : scalar
+        scalar value.
+    B : SparseMatrix
+        Sparse matrix.
+    """
+    raise RuntimeError('power operation between sparse and dense matrix is not supported.')
+
+
+SparseMatrix.__add__ = add
+SparseMatrix.__radd__ = add
+SparseMatrix.__sub__ = sub
+SparseMatrix.__rsub__ = rsub
+SparseMatrix.__mul__ = mul
+SparseMatrix.__rmul__ = mul
+SparseMatrix.__truediv__ = div
+SparseMatrix.__rtruediv__ = rdiv
+SparseMatrix.__pow__ = power
+SparseMatrix.__rpow__ = rpower
\ No newline at end of file
diff --git a/python/dgl/mock_sparse/sp_matrix.py b/python/dgl/mock_sparse/sp_matrix.py
new file mode 100644
index 000000000000..80b1c60e948a
--- /dev/null
+++ b/python/dgl/mock_sparse/sp_matrix.py
@@ -0,0 +1,399 @@
+"""dgl sparse matrix module."""
+from typing import Optional, Tuple
+import torch
+
+class SparseMatrix:
+    r'''Class for sparse matrix.
+
+    Parameters
+    ----------
+    row : tensor
+        The row indices of shape nnz.
+    col : tensor
+        The column indices of shape nnz.
+    val : tensor, optional
+        The values of shape (nnz, *). If None, it will be a tensor of shape (nnz)
+        filled by 1.
+
+    Examples
+    --------
+    Case1: Sparse matrix with row indices, col indices and values (scalar).
+
+    >>> src = torch.tensor([1, 1, 2])
+    >>> dst = torch.tensor([2, 4, 3])
+    >>> val = torch.tensor([1, 1, 1])
+    >>> A = SparseMatrix(src, dst, val)
+    >>> A.shape
+    (3,5)
+    >>> A.row
+    tensor([1, 1, 2])
+    >>> A.val
+    tensor([1., 1., 1.])
+    >>> A.nnz
+    3
+
+    Case2: Sparse matrix with row indices, col indices and values (vector).
+
+    >>> ...
+    >>> val = torch.tensor([[1, 1], [2, 2], [3, 3]])
+    >>> A = SparseMatrix(src, dst, val)
+    >>> A.val
+    tensor([[1, 1],
+            [2, 2],
+            [3, 3]])
+    '''
+    def __init__(self,
+        row: torch.Tensor,
+        col: torch.Tensor,
+        val: Optional[torch.Tensor] = None
+    ):
+        self._row = row
+        self._col = col
+
+        if val is None:
+            val = torch.ones(row.shape[0])
+        self._val = val
+
+        i = torch.cat((row.unsqueeze(0), col.unsqueeze(0)), 0)
+        self.adj = torch.sparse_coo_tensor(i, val).coalesce()
+
+    def __repr__(self):
+        return f'SparseMatrix(indices={self.indices("COO")}, \nvalues={self.val}, \
+                \nshape={self.shape}, nnz={self.nnz})'
+
+    @property
+    def shape(self) -> Tuple[int, ...]:
+        """Shape of the sparse matrix.
+
+        Returns
+        -------
+        tuple[int]
+            The shape of the matrix
+        """
+        return (self.adj.shape[0], self.adj.shape[1])
+
+    @property
+    def nnz(self) -> int:
+        """The number of nonzero elements of the sparse matrix.
+
+        Returns
+        -------
+        int
+            The number of nonzero elements of the matrix
+        """
+        return self.adj._nnz()
+
+    @property
+    def dtype(self) -> torch.dtype:
+        """Data type of the values of the sparse matrix.
+
+        Returns
+        -------
+        torch.dtype
+            Data type of the values of the matrix
+        """
+        return self.adj.dtype
+
+    @property
+    def device(self) -> torch.device:
+        """Device of the sparse matrix.
+
+        Returns
+        -------
+        torch.device
+            Device of the matrix
+        """
+        return self.adj.device
+
+    @property
+    def row(self) -> torch.tensor:
+        """Get the row indices of the nonzero elements.
+
+        Returns
+        -------
+        tensor
+            Row indices of the nonzero elements
+        """
+        return self.adj.indices()[0]
+
+    @property
+    def col(self) -> torch.tensor:
+        """Get the column indices of the nonzero elements.
+
+        Returns
+        -------
+        tensor
+            Column indices of the nonzero elements
+        """
+        return self.adj.indices()[1]
+
+    @property
+    def val(self) -> torch.tensor:
+        """Get the values of the nonzero elements.
+
+        Returns
+        -------
+        tensor
+            Values of the nonzero elements
+        """
+        return self._val
+
+    @val.setter
+    def val(self, x) -> torch.tensor:
+        """Set the values of the nonzero elements."""
+        assert len(x) == self.nnz
+        self._val = x
+
+    def __call__(self, x):
+        """Create a new sparse matrix with the same sparsity as self but different values.
+
+        Parameters
+        ----------
+        x : tensor
+            Values of the new sparse matrix
+
+        Returns
+        -------
+        Class object
+            A new sparse matrix object of the SparseMatrix class
+
+        """
+        assert len(x) == self.nnz
+        return SparseMatrix(self.row, self.col, x)
+
+    def indices(self, format, return_shuffle=False) -> Tuple[torch.tensor, ...]:
+        """Get the indices of the nonzero elements.
+
+        Parameters
+        ----------
+        format : str
+            Sparse matrix storage format. Can be COO or CSR or CSC.
+        return_shuffle: bool
+            If true, return an extra array of the nonzero value IDs
+
+        Returns
+        -------
+        tensor
+            Indices of the nonzero elements
+        """
+        if format == 'COO' and not return_shuffle:
+            return self.adj.indices()
+        else:
+            raise NotImplementedError
+
+    def coo(self) -> Tuple[torch.tensor, ...]:
+        """Get the coordinate (COO) representation of the sparse matrix.
+
+        Returns
+        -------
+        tensor
+            A tensor containing indices and value tensors.
+        """
+        return self
+
+    def csr(self) -> Tuple[torch.tensor, ...]:
+        """Get the CSR (Compressed Sparse Row) representation of the sparse matrix.
+
+        Returns
+        -------
+        tensor
+            A tensor containing compressed row pointers, column indices and value tensors.
+        """
+        return self
+
+    def csc(self) -> Tuple[torch.tensor, ...]:
+        """Get the CSC (Compressed Sparse Column) representation of the sparse matrix.
+
+        Returns
+        -------
+        tensor
+            A tensor containing compressed column pointers, row indices and value tensors.
+        """
+        return self
+
+    def dense(self) -> torch.tensor:
+        """Get the dense representation of the sparse matrix.
+
+        Returns
+        -------
+        tensor
+            Dense representation of the sparse matrix.
+        """
+        return self.adj.to_dense()
+
+def create_from_coo(row: torch.Tensor,
+                    col: torch.Tensor,
+                    val: Optional[torch.Tensor] = None) -> SparseMatrix:
+    """Create a sparse matrix from row and column coordinates.
+
+    Parameters
+    ----------
+    row : tensor
+        The row indices of shape nnz.
+    col : tensor
+        The column indices of shape nnz.
+    val : tensor, optional
+        The values of shape (nnz) or (nnz, D). If None, it will be a tensor of shape (nnz)
+        filled by 1.
+
+    Returns
+    -------
+    SparseMatrix
+        Sparse matrix
+
+    Examples
+    --------
+
+    Case1: Sparse matrix with row and column indices without values.
+
+    >>> src = torch.tensor([1, 1, 2])
+    >>> dst = torch.tensor([2, 4, 3])
+    >>> A = create_from_coo(src, dst)
+    >>> A.shape
+    (3, 5)
+    >>> A.row
+    tensor([1, 1, 2])
+    >>> A.val
+    tensor([1., 1., 1.])
+    >>> A.nnz
+    3
+
+    Case2: Sparse matrix with scalar/vector values. Following example is with
+    vector data.
+
+    >>> val = torch.tensor([[1, 1], [2, 2], [3, 3]])
+    >>> A = create_from_coo(src, dst, val)
+    >>> A.val
+    tensor([[1, 1],
+            [2, 2],
+            [3, 3]])
+    """
+    return SparseMatrix(row, col, val)
+
+def create_from_csr(indptr: torch.Tensor,
+                    indices: torch.Tensor,
+                    val: Optional[torch.Tensor] = None) -> SparseMatrix:
+    """Create a sparse matrix from CSR indices.
+
+    For row i of the sparse matrix
+
+    - the column indices of the nonzero entries are stored in ``indices[indptr[i]: indptr[i+1]]``
+    - the corresponding values are stored in ``val[indptr[i]: indptr[i+1]]``
+
+    Parameters
+    ----------
+    indptr : tensor
+        Pointer to the column indices of shape N + 1, where N is the number of rows.
+    indices : tensor
+        The column indices of shape nnz.
+    val : tensor, optional
+        The values of shape (nnz) or (nnz, D). If None, it will be a tensor of shape (nnz)
+        filled by 1.
+
+    Returns
+    -------
+    SparseMatrix
+        Sparse matrix
+
+    Examples
+    --------
+
+    Case1: Sparse matrix without values
+
+    [[0, 1, 0],
+     [0, 0, 1],
+     [1, 1, 1]]
+
+    >>> indptr = torch.tensor([0, 1, 2, 5])
+    >>> indices = torch.tensor([1, 2, 0, 1, 2])
+    >>> A = create_from_csr(indptr, indices)
+    >>> A.shape
+    (3, 3)
+    >>> A.row
+    tensor([0, 1, 2, 2, 2])
+    >>> A.val
+    tensor([1., 1., 1., 1., 1.])
+    >>> A.nnz
+    5
+
+    Case2: Sparse matrix with scalar/vector values. Following example is with
+    vector data.
+
+    >>> val = torch.tensor([[1, 1], [2, 2], [3, 3], [4, 4], [5, 5]])
+    >>> A = create_from_csr(indptr, indices, val)
+    >>> A.val
+    tensor([[1, 1],
+            [2, 2],
+            [3, 3],
+            [4, 4],
+            [5, 5]])
+    """
+    adj_csr = torch.sparse_csr_tensor(indptr, indices, torch.ones(indices.shape[0]))
+    adj_coo = adj_csr.to_sparse_coo().coalesce()
+    row, col = adj_coo.indices()
+
+    return SparseMatrix(row, col, val)
+
+def create_from_csc(indptr: torch.Tensor,
+                    indices: torch.Tensor,
+                    val: Optional[torch.Tensor] = None) -> SparseMatrix:
+    """Create a sparse matrix from CSC indices.
+
+    For column i of the sparse matrix
+
+    - the row indices of the nonzero entries are stored in ``indices[indptr[i]: indptr[i+1]]``
+    - the corresponding values are stored in ``val[indptr[i]: indptr[i+1]]``
+
+    Parameters
+    ----------
+    indptr : tensor
+        Pointer to the row indices of shape N + 1, where N is the number of columns.
+    indices : tensor
+        The row indices of shape nnz.
+    val : tensor, optional
+        The values of shape (nnz) or (nnz, D). If None, it will be a tensor of shape (nnz)
+        filled by 1.
+
+    Returns
+    -------
+    SparseMatrix
+        Sparse matrix
+
+    Examples
+    --------
+
+    Case1: Sparse matrix without values
+
+    [[0, 1, 0],
+     [0, 0, 1],
+     [1, 1, 1]]
+
+    >>> indptr = torch.tensor([0, 1, 3, 5])
+    >>> indices = torch.tensor([2, 0, 2, 1, 2])
+    >>> A = create_from_csc(indptr, indices)
+    >>> A.shape
+    (3, 3)
+    >>> A.row
+    tensor([0, 1, 2, 2, 2])
+    >>> A.val
+    tensor([1., 1., 1., 1., 1.])
+    >>> A.nnz
+    5
+
+    Case2: Sparse matrix with scalar/vector values. Following example is with
+    vector data.
+
+    >>> val = torch.tensor([[1, 1], [2, 2], [3, 3], [4, 4], [5, 5]])
+    >>> A = create_from_csc(indptr, indices, val)
+    >>> A.val
+    tensor([[1, 1],
+            [2, 2],
+            [3, 3],
+            [4, 4],
+            [5, 5]])
+    """
+    adj_csr = torch.sparse_csr_tensor(indptr, indices, torch.ones(indices.shape[0]))
+    adj_coo = adj_csr.to_sparse_coo().coalesce()
+    col, row = adj_coo.indices()
+
+    return SparseMatrix(row, col, val)
\ No newline at end of file
diff --git a/tests/mock_sparse/test_elementwise_op_sp.py b/tests/mock_sparse/test_elementwise_op_sp.py
new file mode 100644
index 000000000000..bbc724a9d1ef
--- /dev/null
+++ b/tests/mock_sparse/test_elementwise_op_sp.py
@@ -0,0 +1,76 @@
+import numpy as np
+import pytest
+import dgl
+import dgl.backend as F
+import torch
+import numpy
+import operator
+from dgl.mock_dgl_sparse.python import elementwise_sparse
+from dgl.mock_dgl_sparse.python.sp_matrix import SparseMatrix
+parametrize_idtype = pytest.mark.parametrize("idtype", [F.int32, F.int64])
+parametrize_dtype = pytest.mark.parametrize('dtype', [F.float32, F.float64])
+
+
+ops = {
+    "+": operator.add,
+    "-": operator.sub,
+    "*": operator.mul,
+    "/": operator.truediv,
+}
+
+@parametrize_idtype
+@parametrize_dtype
+def test_sparse_op_sparse(idtype, dtype):
+    row = torch.randint(1, 500, (100,))
+    col = torch.randint(1, 500, (100,))
+    val = torch.rand(100)
+    A = SparseMatrix(row, col, val)
+
+    def _test(A, op):
+        out_torch = eval("A.val " + op + " A.val")
+        out_spMat = eval("A " + op + " A").val
+        assert np.allclose(out_torch, out_spMat, rtol=1e-4, atol=1e-4)
+        w1 = torch.rand(100)
+        w2 = torch.rand(100)
+        out_torch2 = eval("w1 " + op + " w2")
+        out_spMat2 = eval("A(w1) " + op + " A(w2)").val
+        assert np.allclose(out_torch2, out_spMat2, rtol=1e-4, atol=1e-4)
+
+    for op in ops:
+        _test(A, op)
+
+@parametrize_idtype
+@parametrize_dtype
+@pytest.mark.parametrize('v_scalar', [2, 2.5])
+def test_sparse_op_scalar(idtype, dtype, v_scalar):
+    row = torch.randint(1, 500, (100,))
+    col = torch.randint(1, 500, (100,))
+    val = torch.rand(100)
+    A = SparseMatrix(row, col, val)
+    out_torch = A.val * v_scalar
+    out_spMat = A * v_scalar
+    assert np.allclose(out_torch, out_spMat.val, rtol=1e-4, atol=1e-4)
+    out_torch = A.val / v_scalar
+    out_spMat = A / v_scalar
+    assert np.allclose(out_torch, out_spMat.val, rtol=1e-4, atol=1e-4)
+    out_torch = pow(A.val, v_scalar)
+    out_spMat = pow(A, v_scalar)
+    assert np.allclose(out_torch, out_spMat.val, rtol=1e-4, atol=1e-4)
+
+@parametrize_idtype
+@parametrize_dtype
+@pytest.mark.parametrize('v_scalar', [2, 2.5])
+def test_scalar_op_sparse(idtype, dtype, v_scalar):
+    row = torch.randint(1, 500, (100,))
+    col = torch.randint(1, 500, (100,))
+    val = torch.rand(100)
+    A = SparseMatrix(row, col, val)
+    out_torch = v_scalar * A.val
+    out_spMat = v_scalar * A
+    assert np.allclose(out_torch, out_spMat.val, rtol=1e-4, atol=1e-4)
+
+if __name__ == '__main__':
+    test_sparse_op_sparse()
+    test_sparse_op_scalar()
+    test_scalar_op_sparse()
+


Code review comments:
diff_hunk: @@ -0,0 +1,76 @@
+import numpy as np
+import pytest
+import dgl
+import dgl.backend as F
+import torch
+import numpy
+import operator
+from dgl.mock_dgl_sparse.python import elementwise_sparse
+from dgl.mock_dgl_sparse.python.sp_matrix import SparseMatrix
+parametrize_idtype = pytest.mark.parametrize("idtype", [F.int32, F.int64])
+parametrize_dtype = pytest.mark.parametrize('dtype', [F.float32, F.float64])
+
+
+ops = {
+    "+": operator.add,
+    "-": operator.sub,
+    "*": operator.mul,
+    "/": operator.truediv,
+}
+
+@parametrize_idtype
+@parametrize_dtype
+def test_sparse_op_sparse(idtype, dtype):
path: tests/mock_sparse/test_elementwise_op_sp.py
text: @user1:
Add a new argument op and use `pytest.mark.parametrize` to test different elementwise operators. This will save the for loop at L39.

@author:
Done.

diff_hunk: @@ -0,0 +1,219 @@
+"""dgl elementwise operators for sparse matrix module."""
+import torch
+from .sp_matrix import SparseMatrix
+
+def check_sparsity(A, B):
+    """Checks if two sparse matrices have the same sparsity structure."""
+    if not torch.equal(A.indices("COO"), B.indices("COO")):
+        raise ValueError('The two input matrices have different sparsity pattern')
+
+def add(A, B):
+    """Elementwise addition.
+
+    Parameters
+    ----------
+    A : SparseMatrix
+        Sparse matrix
+    B : SparseMatrix
+        Sparse matrix
+
+    Returns
+    -------
+    SparseMatrix
+        Sparse matrix
+
+    Examples
+    --------
+    >>> row = torch.tensor([1, 1, 2])
+    >>> col = torch.tensor([2, 4, 3])
+    >>> w1 = torch.tensor([3 ,4, 5])
+    >>> w2 = torch.tensor([-1, -3, -3])
+    >>> A = SparseMatrix(row, col, w1)
+    >>> C = A + A
+    >>> C.val
+    tensor([ 6,  8, 10])
+    >>> C = A(w2) + A(w1)
+    tensor([2, 1, 2])
+    """
+    if isinstance(A, SparseMatrix) and isinstance(B, SparseMatrix):
+        assert A.shape == B.shape, 'The shape of sparse matrix A {} and' \
+        'B {} must match'.format(A.shape, B.shape)
+        check_sparsity(A, B)
+        return SparseMatrix(A.row, A.col, A.val + B.val)
+    raise RuntimeError('Elementwise add between sparse and dense matrix is not supported.')
path: python/dgl/mock_sparse/elementwise_op_sp.py
text: @user1:
Is there any cases that will make this error message inaccurate? e.g., users pass in a scalar `A` and a sparse `B`. One way to tackle this is to have a more general error message, e.g., 'Elementwise addition between {type(A)} and {type(B)} are not supported.'

diff_hunk: @@ -0,0 +1,219 @@
+"""dgl elementwise operators for sparse matrix module."""
+import torch
+from .sp_matrix import SparseMatrix
+
+def check_sparsity(A, B):
+    """Checks if two sparse matrices have the same sparsity structure."""
+    if not torch.equal(A.indices("COO"), B.indices("COO")):
+        raise ValueError('The two input matrices have different sparsity pattern')
+
+def add(A, B):
+    """Elementwise addition.
+
+    Parameters
+    ----------
+    A : SparseMatrix
+        Sparse matrix
+    B : SparseMatrix
+        Sparse matrix
+
+    Returns
+    -------
+    SparseMatrix
+        Sparse matrix
+
+    Examples
+    --------
+    >>> row = torch.tensor([1, 1, 2])
+    >>> col = torch.tensor([2, 4, 3])
+    >>> w1 = torch.tensor([3 ,4, 5])
+    >>> w2 = torch.tensor([-1, -3, -3])
+    >>> A = SparseMatrix(row, col, w1)
+    >>> C = A + A
+    >>> C.val
+    tensor([ 6,  8, 10])
+    >>> C = A(w2) + A(w1)
+    tensor([2, 1, 2])
+    """
+    if isinstance(A, SparseMatrix) and isinstance(B, SparseMatrix):
+        assert A.shape == B.shape, 'The shape of sparse matrix A {} and' \
+        'B {} must match'.format(A.shape, B.shape)
+        check_sparsity(A, B)
+        return SparseMatrix(A.row, A.col, A.val + B.val)
+    raise RuntimeError('Elementwise add between sparse and dense matrix is not supported.')
+
+def sub(A, B):
+    """Elementwise subtraction.
+
+    Parameters
+    ----------
+    A : SparseMatrix
+        Sparse matrix
+    B : SparseMatrix
+        Sparse matrix
+
+    Returns
+    -------
+    SparseMatrix
+        Sparse matrix
+
+    Examples
+    --------
+    ...
+    >>> C = A - A
+    >>> C.val
+    tensor([0, 0, 0])
+    >>> C = A(w2) - A(w1)
+    tensor([-4, -7, -8])
+    """
+    if isinstance(A, SparseMatrix) and isinstance(B, SparseMatrix):
+        assert A.shape == B.shape, 'The shape of sparse matrix A {} and' \
+        'B {} must match'.format(A.shape, B.shape)
+        check_sparsity(A, B)
path: python/dgl/mock_sparse/elementwise_op_sp.py
text: @user1:
Add a TODO that we will support adding two matrices of different sparsity later.

@author:
Added support to add two matrices of different sparsity.

diff_hunk: @@ -0,0 +1,76 @@
+import numpy as np
+import pytest
+import dgl
+import dgl.backend as F
+import torch
+import numpy
+import operator
+from dgl.mock_dgl_sparse.python import elementwise_sparse
+from dgl.mock_dgl_sparse.python.sp_matrix import SparseMatrix
+parametrize_idtype = pytest.mark.parametrize("idtype", [F.int32, F.int64])
+parametrize_dtype = pytest.mark.parametrize('dtype', [F.float32, F.float64])
+
+
+ops = {
+    "+": operator.add,
+    "-": operator.sub,
+    "*": operator.mul,
+    "/": operator.truediv,
+}
+
+@parametrize_idtype
+@parametrize_dtype
+def test_sparse_op_sparse(idtype, dtype):
+    row = torch.randint(1, 500, (100,))
+    col = torch.randint(1, 500, (100,))
+    val = torch.rand(100)
+    A = SparseMatrix(row, col, val)
+
+    def _test(A, op):
+        out_torch = eval("A.val " + op + " A.val")
path: tests/mock_sparse/test_elementwise_op_sp.py
text: @user1:
directly use `op(A.val, A.val)`. Use `eval` is a bit confusing.

@author:
Done.

diff_hunk: @@ -0,0 +1,399 @@
+"""dgl sparse matrix module."""
+from typing import Optional, Tuple
+import torch
+
+class SparseMatrix:
+    r'''Class for sparse matrix.
+
+    Parameters
+    ----------
+    row : tensor
+        The row indices of shape nnz.
+    col : tensor
+        The column indices of shape nnz.
+    val : tensor, optional
+        The values of shape (nnz, *). If None, it will be a tensor of shape (nnz)
+        filled by 1.
+
+    Examples
+    --------
+    Case1: Sparse matrix with row indices, col indices and values (scalar).
+
+    >>> src = torch.tensor([1, 1, 2])
+    >>> dst = torch.tensor([2, 4, 3])
+    >>> val = torch.tensor([1, 1, 1])
+    >>> A = SparseMatrix(src, dst, val)
+    >>> A.shape
+    (3,5)
+    >>> A.row
+    tensor([1, 1, 2])
+    >>> A.val
+    tensor([1., 1., 1.])
+    >>> A.nnz
+    3
+
+    Case2: Sparse matrix with row indices, col indices and values (vector).
+
+    >>> ...
+    >>> val = torch.tensor([[1, 1], [2, 2], [3, 3]])
+    >>> A = SparseMatrix(src, dst, val)
+    >>> A.val
+    tensor([[1, 1],
+            [2, 2],
+            [3, 3]])
+    '''
+    def __init__(self,
+        row: torch.Tensor,
+        col: torch.Tensor,
+        val: Optional[torch.Tensor] = None
+    ):
+        self._row = row
+        self._col = col
+
+        if val is None:
+            val = torch.ones(row.shape[0])
+        self._val = val
+
+        i = torch.cat((row.unsqueeze(0), col.unsqueeze(0)), 0)
+        self.adj = torch.sparse_coo_tensor(i, val).coalesce()
+
+    def __repr__(self):
+        return f'SparseMatrix(indices={self.indices("COO")}, \nvalues={self.val}, \
+                \nshape={self.shape}, nnz={self.nnz})'
+
+    @property
+    def shape(self) -> Tuple[int, ...]:
+        """Shape of the sparse matrix.
+
+        Returns
+        -------
+        tuple[int]
+            The shape of the matrix
+        """
+        return (self.adj.shape[0], self.adj.shape[1])
+
+    @property
+    def nnz(self) -> int:
+        """The number of nonzero elements of the sparse matrix.
+
+        Returns
+        -------
+        int
+            The number of nonzero elements of the matrix
+        """
+        return self.adj._nnz()
+
+    @property
+    def dtype(self) -> torch.dtype:
+        """Data type of the values of the sparse matrix.
+
+        Returns
+        -------
+        torch.dtype
+            Data type of the values of the matrix
+        """
+        return self.adj.dtype
+
+    @property
+    def device(self) -> torch.device:
+        """Device of the sparse matrix.
+
+        Returns
+        -------
+        torch.device
+            Device of the matrix
+        """
+        return self.adj.device
+
+    @property
+    def row(self) -> torch.tensor:
+        """Get the row indices of the nonzero elements.
+
+        Returns
+        -------
+        tensor
+            Row indices of the nonzero elements
+        """
+        return self.adj.indices()[0]
+
+    @property
+    def col(self) -> torch.tensor:
+        """Get the column indices of the nonzero elements.
+
+        Returns
+        -------
+        tensor
+            Column indices of the nonzero elements
+        """
+        return self.adj.indices()[1]
+
+    @property
+    def val(self) -> torch.tensor:
+        """Get the values of the nonzero elements.
+
+        Returns
+        -------
+        tensor
+            Values of the nonzero elements
+        """
+        return self._val
+
+    @val.setter
+    def val(self, x) -> torch.tensor:
+        """Set the values of the nonzero elements."""
+        assert len(x) == self.nnz
+        self._val = x
+
+    def __call__(self, x):
+        """Create a new sparse matrix with the same sparsity as self but different values.
+
+        Parameters
+        ----------
+        x : tensor
+            Values of the new sparse matrix
+
+        Returns
+        -------
+        Class object
+            A new sparse matrix object of the SparseMatrix class
+
+        """
+        assert len(x) == self.nnz
+        return SparseMatrix(self.row, self.col, x)
+
+    def indices(self, format, return_shuffle=False) -> Tuple[torch.tensor, ...]:
+        """Get the indices of the nonzero elements.
+
+        Parameters
+        ----------
+        format : str
+            Sparse matrix storage format. Can be COO or CSR or CSC.
+        return_shuffle: bool
+            If true, return an extra array of the nonzero value IDs
+
+        Returns
+        -------
+        tensor
+            Indices of the nonzero elements
+        """
+        if format == 'COO' and not return_shuffle:
+            return self.adj.indices()
+        else:
+            raise NotImplementedError
+
+    def coo(self) -> Tuple[torch.tensor, ...]:
+        """Get the coordinate (COO) representation of the sparse matrix.
+
+        Returns
+        -------
+        tensor
+            A tensor containing indices and value tensors.
+        """
+        return self
+
+    def csr(self) -> Tuple[torch.tensor, ...]:
+        """Get the CSR (Compressed Sparse Row) representation of the sparse matrix.
+
+        Returns
+        -------
+        tensor
+            A tensor containing compressed row pointers, column indices and value tensors.
+        """
+        return self
+
+    def csc(self) -> Tuple[torch.tensor, ...]:
+        """Get the CSC (Compressed Sparse Column) representation of the sparse matrix.
+
+        Returns
+        -------
+        tensor
+            A tensor containing compressed column pointers, row indices and value tensors.
+        """
+        return self
+
+    def dense(self) -> torch.tensor:
+        """Get the dense representation of the sparse matrix.
+
+        Returns
+        -------
+        tensor
+            Dense representation of the sparse matrix.
+        """
+        return self.adj.to_dense()
+
+def create_from_coo(row: torch.Tensor,
+                    col: torch.Tensor,
+                    val: Optional[torch.Tensor] = None) -> SparseMatrix:
+    """Create a sparse matrix from row and column coordinates.
+
+    Parameters
+    ----------
+    row : tensor
+        The row indices of shape nnz.
+    col : tensor
+        The column indices of shape nnz.
+    val : tensor, optional
+        The values of shape (nnz) or (nnz, D). If None, it will be a tensor of shape (nnz)
+        filled by 1.
+
+    Returns
+    -------
+    SparseMatrix
+        Sparse matrix
+
+    Examples
+    --------
+
+    Case1: Sparse matrix with row and column indices without values.
+
+    >>> src = torch.tensor([1, 1, 2])
+    >>> dst = torch.tensor([2, 4, 3])
+    >>> A = create_from_coo(src, dst)
+    >>> A.shape
path: python/dgl/mock_sparse/sp_matrix.py
text: @user1:
Replace L252-L259 with print(A) and its corresponding output

@author:
Done.

diff_hunk: @@ -0,0 +1,219 @@
+"""dgl elementwise operators for sparse matrix module."""
+import torch
+from .sp_matrix import SparseMatrix
+
+def check_sparsity(A, B):
+    """Checks if two sparse matrices have the same sparsity structure."""
+    if not torch.equal(A.indices("COO"), B.indices("COO")):
+        raise ValueError('The two input matrices have different sparsity pattern')
+
+def add(A, B):
+    """Elementwise addition.
+
+    Parameters
+    ----------
+    A : SparseMatrix
+        Sparse matrix
+    B : SparseMatrix
+        Sparse matrix
+
+    Returns
+    -------
+    SparseMatrix
+        Sparse matrix
+
+    Examples
+    --------
+    >>> row = torch.tensor([1, 1, 2])
+    >>> col = torch.tensor([2, 4, 3])
+    >>> w1 = torch.tensor([3 ,4, 5])
+    >>> w2 = torch.tensor([-1, -3, -3])
+    >>> A = SparseMatrix(row, col, w1)
+    >>> C = A + A
+    >>> C.val
+    tensor([ 6,  8, 10])
+    >>> C = A(w2) + A(w1)
+    tensor([2, 1, 2])
+    """
+    if isinstance(A, SparseMatrix) and isinstance(B, SparseMatrix):
+        assert A.shape == B.shape, 'The shape of sparse matrix A {} and' \
+        'B {} must match'.format(A.shape, B.shape)
+        check_sparsity(A, B)
+        return SparseMatrix(A.row, A.col, A.val + B.val)
+    raise RuntimeError('Elementwise add between sparse and dense matrix is not supported.')
+
+def sub(A, B):
+    """Elementwise subtraction.
+
+    Parameters
+    ----------
+    A : SparseMatrix
+        Sparse matrix
+    B : SparseMatrix
+        Sparse matrix
+
+    Returns
+    -------
+    SparseMatrix
+        Sparse matrix
+
+    Examples
+    --------
+    ...
path: python/dgl/mock_sparse/elementwise_op_sp.py
text: @user1:
Please make each example self-contained.

diff_hunk: @@ -0,0 +1,76 @@
+import numpy as np
+import pytest
+import dgl
+import dgl.backend as F
+import torch
+import numpy
+import operator
+from dgl.mock_dgl_sparse.python import elementwise_sparse
+from dgl.mock_dgl_sparse.python.sp_matrix import SparseMatrix
+parametrize_idtype = pytest.mark.parametrize("idtype", [F.int32, F.int64])
+parametrize_dtype = pytest.mark.parametrize('dtype', [F.float32, F.float64])
+
+
+ops = {
+    "+": operator.add,
+    "-": operator.sub,
+    "*": operator.mul,
+    "/": operator.truediv,
+}
+
+@parametrize_idtype
+@parametrize_dtype
+def test_sparse_op_sparse(idtype, dtype):
+    row = torch.randint(1, 500, (100,))
+    col = torch.randint(1, 500, (100,))
+    val = torch.rand(100)
+    A = SparseMatrix(row, col, val)
+
+    def _test(A, op):
+        out_torch = eval("A.val " + op + " A.val")
+        out_spMat = eval("A " + op + " A").val
+        assert np.allclose(out_torch, out_spMat, rtol=1e-4, atol=1e-4)
+        w1 = torch.rand(100)
+        w2 = torch.rand(100)
+        out_torch2 = eval("w1 " + op + " w2")
+        out_spMat2 = eval("A(w1) " + op + " A(w2)").val
+        assert np.allclose(out_torch2, out_spMat2, rtol=1e-4, atol=1e-4)
+
+    for op in ops:
+        _test(A, op)
+
+@parametrize_idtype
+@parametrize_dtype
+@pytest.mark.parametrize('v_scalar', [2, 2.5])
+def test_sparse_op_scalar(idtype, dtype, v_scalar):
+    row = torch.randint(1, 500, (100,))
+    col = torch.randint(1, 500, (100,))
+    val = torch.rand(100)
+    A = SparseMatrix(row, col, val)
+    out_torch = A.val * v_scalar
+    out_spMat = A * v_scalar
+    assert np.allclose(out_torch, out_spMat.val, rtol=1e-4, atol=1e-4)
+    out_torch = A.val / v_scalar
+    out_spMat = A / v_scalar
+    assert np.allclose(out_torch, out_spMat.val, rtol=1e-4, atol=1e-4)
+    out_torch = pow(A.val, v_scalar)
+    out_spMat = pow(A, v_scalar)
+    assert np.allclose(out_torch, out_spMat.val, rtol=1e-4, atol=1e-4)
+
+@parametrize_idtype
+@parametrize_dtype
+@pytest.mark.parametrize('v_scalar', [2, 2.5])
+def test_scalar_op_sparse(idtype, dtype, v_scalar):
+    row = torch.randint(1, 500, (100,))
+    col = torch.randint(1, 500, (100,))
+    val = torch.rand(100)
+    A = SparseMatrix(row, col, val)
+    out_torch = v_scalar * A.val
+    out_spMat = v_scalar * A
path: tests/mock_sparse/test_elementwise_op_sp.py
text: @user1:
Missing `A / v_scalar`

diff_hunk: @@ -0,0 +1,76 @@
+import numpy as np
path: tests/mock_sparse/test_elementwise_op_sp.py
text: @user1:
You need to put the file under `tests/pytorch`. Otherwise, the tests won't be run.