[Sparse] Update doc string of element-wise ops for DiagMatrix and SparseMatrix. (#5136)

* doc

* minor

* black

* no-lint

* reorg
Co-authored-by: Steve <ubuntu@ip-172-31-34-29.ap-northeast-1.compute.internal>

docs/source/api/python/dgl.sparse_v0.rst

@@ -153,6 +153,19 @@ Operators
 ---------
 .. currentmodule:: dgl.sparse
 
+Elementwise Operators
+````````
+
+.. autosummary::
+    :toctree: ../../generated/
+
+    add
+    sub
+    power
+
+Matrix Multiplication
+````````
+
 .. autosummary::
     :toctree: ../../generated/
 
@@ -162,4 +175,11 @@ Operators
     mm
     sddmm
     bsddmm
+
+Non-linear activation functions
+````````
+
+.. autosummary::
+    :toctree: ../../generated/
+
     softmax

python/dgl/sparse/elementwise_op.py

+# pylint: disable=anomalous-backslash-in-string
 """DGL elementwise operator module."""
 from typing import Union
 
@@ -10,7 +11,18 @@ __all__ = ["add", "power"]
 def add(
     A: Union[DiagMatrix, SparseMatrix], B: Union[DiagMatrix, SparseMatrix]
 ) -> Union[DiagMatrix, SparseMatrix]:
-    """Elementwise addition
+    """Elementwise additions for `DiagMatrix` and `SparseMatrix`.
+
+    The supported combinations are shown as follow.
+    +--------------+------------+--------------+--------+
+    |     A \ B    | DiagMatrix | SparseMatrix | scalar |
+    +--------------+------------+--------------+--------+
+    |  DiagMatrix  |     ✅     |      ✅      |   🚫   |
+    +--------------+------------+--------------+--------+
+    | SparseMatrix |     ✅     |      ✅      |   🚫   |
+    +--------------+------------+--------------+--------+
+    |    scalar    |     🚫     |      🚫      |   🚫   |
+    +--------------+------------+--------------+--------+
 
     Parameters
     ----------
@@ -34,18 +46,65 @@ def add(
     >>> B = diag(torch.arange(1, 4))
     >>> A + B
     SparseMatrix(indices=tensor([[0, 0, 1, 1, 2],
-            [0, 1, 0, 1, 2]]),
+                                 [0, 1, 0, 1, 2]]),
     values=tensor([ 1, 20, 10,  2, 33]),
     shape=(3, 3), nnz=5)
     """
     return A + B
 
 
+def sub(A: Union[DiagMatrix], B: Union[DiagMatrix]) -> Union[DiagMatrix]:
+    """Elementwise subtraction for `DiagMatrix` and `SparseMatrix`.
+
+    The supported combinations are shown as follow.
+    +--------------+------------+--------------+--------+
+    |     A \ B    | DiagMatrix | SparseMatrix | scalar |
+    +--------------+------------+--------------+--------+
+    |  DiagMatrix  |     ✅     |      🚫      |   🚫   |
+    +--------------+------------+--------------+--------+
+    | SparseMatrix |     🚫     |      🚫      |   🚫   |
+    +--------------+------------+--------------+--------+
+    |    scalar    |     🚫     |      🚫      |   🚫   |
+    +--------------+------------+--------------+--------+
+
+    Parameters
+    ----------
+    A : DiagMatrix
+        Diagonal matrix
+    B : DiagMatrix
+        Diagonal matrix
+
+    Returns
+    -------
+    DiagMatrix
+        Diagonal matrix
+
+    Examples
+    --------
+    >>> A = diag(torch.arange(1, 4))
+    >>> B = diag(torch.arange(10, 13))
+    >>> A - B
+    DiagMatrix(val=tensor([-9, -9, -9]),
+    shape=(3, 3))
+    """
+    return A - B
+
+
 def power(
     A: Union[SparseMatrix, DiagMatrix], scalar: Union[float, int]
 ) -> Union[SparseMatrix, DiagMatrix]:
-    """Take the power of each nonzero element and return a matrix with
-    the result.
+    """Elementwise exponentiation for `DiagMatrix` and `SparseMatrix`.
+
+    The supported combinations are shown as follow.
+    +--------------+------------+--------------+--------+
+    |     A \ B    | DiagMatrix | SparseMatrix | scalar |
+    +--------------+------------+--------------+--------+
+    |  DiagMatrix  |     🚫     |      🚫      |   ✅   |
+    +--------------+------------+--------------+--------+
+    | SparseMatrix |     🚫     |      🚫      |   ✅   |
+    +--------------+------------+--------------+--------+
+    |    scalar    |     🚫     |      🚫      |   🚫   |
+    +--------------+------------+--------------+--------+
 
     Parameters
     ----------
@@ -68,7 +127,7 @@ def power(
     >>> A = from_coo(row, col, val)
     >>> power(A, 2)
     SparseMatrix(indices=tensor([[1, 0, 2],
-            [0, 3, 2]]),
+                                 [0, 3, 2]]),
     values=tensor([100, 400, 900]),
     shape=(3, 4), nnz=3)