[Sparse] neg and inv (#4980)

* Update

* update
Co-authored-by: Ubuntu <ubuntu@ip-172-31-36-188.ap-northeast-1.compute.internal>

python/dgl/mock_sparse2/__init__.py

 """dgl sparse class."""
-import sys
 import os
+import sys
+
 import torch
 
 from .._ffi import libinfo
-from .sparse_matrix import *
 from .diag_matrix import *
 from .elementwise_op import *
+from .sparse_matrix import *
+from .unary_op_diag import *
+from .unary_op_sp import *
 
 
 def load_dgl_sparse():

python/dgl/mock_sparse2/unary_op_diag.py

+"""DGL unary operators for diagonal matrix module."""
+# pylint: disable= invalid-name
+from .diag_matrix import DiagMatrix, diag
+
+
+def neg(D: DiagMatrix) -> DiagMatrix:
+    """Return a new diagonal matrix with the negation of the original nonzero
+    values.
+
+    Returns
+    -------
+    DiagMatrix
+        Negation of the diagonal matrix
+
+    Examples
+    --------
+
+    >>> val = torch.arange(3).float()
+    >>> mat = diag(val)
+    >>> mat = -mat
+    >>> print(mat)
+    DiagMatrix(val=tensor([-0., -1., -2.]),
+    shape=(3, 3))
+    """
+    return diag(-D.val, D.shape)
+
+
+def inv(D: DiagMatrix) -> DiagMatrix:
+    """Return the inverse of the diagonal matrix.
+
+    This function only supports square matrices with scalar nonzero values.
+
+    Returns
+    -------
+    DiagMatrix
+        Inverse of the diagonal matrix
+
+    Examples
+    --------
+
+    >>> val = torch.arange(1, 4).float()
+    >>> mat = diag(val)
+    >>> mat = mat.inv()
+    >>> print(mat)
+    DiagMatrix(val=tensor([1.0000, 0.5000, 0.3333]),
+    shape=(3, 3))
+    """
+    num_rows, num_cols = D.shape
+    assert num_rows == num_cols, f"Expect a square matrix, got shape {D.shape}"
+    assert len(D.val.shape) == 1, "inv only supports 1D nonzero val"
+
+    return diag(1.0 / D.val, D.shape)
+
+
+DiagMatrix.neg = neg
+DiagMatrix.__neg__ = neg
+DiagMatrix.inv = inv

python/dgl/mock_sparse2/unary_op_sp.py

+"""DGL unary operators for sparse matrix module."""
+from .sparse_matrix import SparseMatrix, val_like
+
+
+def neg(A: SparseMatrix) -> SparseMatrix:
+    """Return a new sparse matrix with the negation of the original nonzero
+    values.
+
+    Returns
+    -------
+    SparseMatrix
+        Negation of the sparse matrix
+
+    Examples
+    --------
+
+    >>> row = torch.tensor([1, 1, 3])
+    >>> col = torch.tensor([1, 2, 3])
+    >>> val = torch.tensor([1., 1., 2.])
+    >>> A = create_from_coo(row, col, val)
+    >>> A = -A
+    >>> print(A)
+    SparseMatrix(indices=tensor([[1, 1, 3],
+        [1, 2, 3]]),
+    values=tensor([-1., -1., -2.]),
+    shape=(4, 4), nnz=3)
+    """
+    return val_like(A, -A.val)
+
+
+SparseMatrix.neg = neg
+SparseMatrix.__neg__ = neg

tests/pytorch/mock_sparse2/test_unary_op_diag.py

+import pytest
+import torch
+import sys
+
+import backend as F
+
+from dgl.mock_sparse2 import diag
+
+
+def test_neg():
+    ctx = F.ctx()
+    val = torch.arange(3).float().to(ctx)
+    D = diag(val)
+    neg_D = -D
+    assert D.shape == neg_D.shape
+    assert torch.allclose(-D.val, neg_D.val)
+    assert D.val.device == neg_D.val.device
+
+
+def test_inv():
+    ctx = F.ctx()
+    val = torch.arange(1, 4).float().to(ctx)
+    D = diag(val)
+    inv_D = D.inv()
+    assert D.shape == inv_D.shape
+    assert torch.allclose(1.0 / D.val, inv_D.val)
+    assert D.val.device == inv_D.val.device

tests/pytorch/mock_sparse2/test_unary_op_sp.py

+import pytest
+import torch
+import sys
+
+import backend as F
+
+from dgl.mock_sparse2 import create_from_coo
+
+# FIXME(issue #4818): Skipping tests on win.
+if not sys.platform.startswith("linux"):
+    pytest.skip("skipping tests on win", allow_module_level=True)
+
+
+def test_neg():
+    ctx = F.ctx()
+    row = torch.tensor([1, 1, 3]).to(ctx)
+    col = torch.tensor([1, 2, 3]).to(ctx)
+    val = torch.tensor([1.0, 1.0, 2.0]).to(ctx)
+    A = create_from_coo(row, col, val)
+    neg_A = -A
+    assert A.shape == neg_A.shape
+    assert A.nnz == neg_A.nnz
+    assert torch.allclose(-A.val, neg_A.val)
+    assert torch.allclose(torch.stack(A.coo()), torch.stack(neg_A.coo()))
+    assert A.val.device == neg_A.val.device