[Sparse] Clarification on dimensions for sparse reduction and softmax docstrings (#5197)

* fix sparse reduce and softmax doc

* Update softmax.py

python/dgl/sparse/reduction.py

@@ -21,11 +21,12 @@ def reduce(input: SparseMatrix, dim: Optional[int] = None, rtype: str = "sum"):
         The input sparse matrix
     dim : int, optional
         The dimension to reduce, must be either 0 (by rows) or 1 (by columns)
-        or None (on all non-zero entries)
+        or None (on both rows and columns simultaneously)
 
-        If :attr:`dim` is None, it reduces all the elements in the sparse
-        matrix. Otherwise, it reduces on the row (``dim=0``) or column
-        (``dim=1``) dimension, producing a tensor of shape
+        If :attr:`dim` is None, it reduces both the rows and the columns
+        in the sparse matrix, producing a tensor of shape
+        ``input.val.shape[1:]``. Otherwise, it reduces on the row (``dim=0``)
+        or column (``dim=1``) dimension, producing a tensor of shape
         ``(input.shape[1],) + input.val.shape[1:]`` or
         ``(input.shape[0],) + input.val.shape[1:]``.
     rtype: str, optional
@@ -93,11 +94,12 @@ def sum(input: SparseMatrix, dim: Optional[int] = None):
         The input sparse matrix
     dim : int, optional
         The dimension to reduce, must be either 0 (by rows) or 1 (by columns)
-        or None (on all non-zero entries)
+        or None (on both rows and columns simultaneously)
 
-        If :attr:`dim` is None, it reduces all the elements in the sparse
-        matrix. Otherwise, it reduces on the row (``dim=0``) or column
-        (``dim=1``) dimension, producing a tensor of shape
+        If :attr:`dim` is None, it reduces both the rows and the columns
+        in the sparse matrix, producing a tensor of shape
+        ``input.val.shape[1:]``. Otherwise, it reduces on the row (``dim=0``)
+        or column (``dim=1``) dimension, producing a tensor of shape
         ``(input.shape[1],) + input.val.shape[1:]`` or
         ``(input.shape[0],) + input.val.shape[1:]``.
 
@@ -151,11 +153,12 @@ def smax(input: SparseMatrix, dim: Optional[int] = None):
         The input sparse matrix
     dim : int, optional
         The dimension to reduce, must be either 0 (by rows) or 1 (by columns)
-        or None (on all non-zero entries)
+        or None (on both rows and columns simultaneously)
 
-        If :attr:`dim` is None, it reduces all the elements in the sparse
-        matrix. Otherwise, it reduces on the row (``dim=0``) or column
-        (``dim=1``) dimension, producing a tensor of shape
+        If :attr:`dim` is None, it reduces both the rows and the columns
+        in the sparse matrix, producing a tensor of shape
+        ``input.val.shape[1:]``. Otherwise, it reduces on the row (``dim=0``)
+        or column (``dim=1``) dimension, producing a tensor of shape
         ``(input.shape[1],) + input.val.shape[1:]`` or
         ``(input.shape[0],) + input.val.shape[1:]``.
 
@@ -210,11 +213,12 @@ def smin(input: SparseMatrix, dim: Optional[int] = None):
         The input sparse matrix
     dim : int, optional
         The dimension to reduce, must be either 0 (by rows) or 1 (by columns)
-        or None (on all non-zero entries)
+        or None (on both rows and columns simultaneously)
 
-        If :attr:`dim` is None, it reduces all the elements in the sparse
-        matrix. Otherwise, it reduces on the row (``dim=0``) or column
-        (``dim=1``) dimension, producing a tensor of shape
+        If :attr:`dim` is None, it reduces both the rows and the columns
+        in the sparse matrix, producing a tensor of shape
+        ``input.val.shape[1:]``. Otherwise, it reduces on the row (``dim=0``)
+        or column (``dim=1``) dimension, producing a tensor of shape
         ``(input.shape[1],) + input.val.shape[1:]`` or
         ``(input.shape[0],) + input.val.shape[1:]``.
 
@@ -273,11 +277,12 @@ def smean(input: SparseMatrix, dim: Optional[int] = None):
         The input sparse matrix
     dim : int, optional
         The dimension to reduce, must be either 0 (by rows) or 1 (by columns)
-        or None (on all non-zero entries)
+        or None (on both rows and columns simultaneously)
 
-        If :attr:`dim` is None, it reduces all the elements in the sparse
-        matrix. Otherwise, it reduces on the row (``dim=0``) or column
-        (``dim=1``) dimension, producing a tensor of shape
+        If :attr:`dim` is None, it reduces both the rows and the columns
+        in the sparse matrix, producing a tensor of shape
+        ``input.val.shape[1:]``. Otherwise, it reduces on the row (``dim=0``)
+        or column (``dim=1``) dimension, producing a tensor of shape
         ``(input.shape[1],) + input.val.shape[1:]`` or
         ``(input.shape[0],) + input.val.shape[1:]``.
 
@@ -336,11 +341,12 @@ def sprod(input: SparseMatrix, dim: Optional[int] = None):
         The input sparse matrix
     dim : int, optional
         The dimension to reduce, must be either 0 (by rows) or 1 (by columns)
-        or None (on all non-zero entries)
+        or None (on both rows and columns simultaneously)
 
-        If :attr:`dim` is None, it reduces all the elements in the sparse
-        matrix. Otherwise, it reduces on the row (``dim=0``) or column
-        (``dim=1``) dimension, producing a tensor of shape
+        If :attr:`dim` is None, it reduces both the rows and the columns
+        in the sparse matrix, producing a tensor of shape
+        ``input.val.shape[1:]``. Otherwise, it reduces on the row (``dim=0``)
+        or column (``dim=1``) dimension, producing a tensor of shape
         ``(input.shape[1],) + input.val.shape[1:]`` or
         ``(input.shape[0],) + input.val.shape[1:]``.
 

python/dgl/sparse/softmax.py

@@ -9,7 +9,10 @@ __all__ = ["softmax"]
 
 
 def softmax(input: SparseMatrix) -> SparseMatrix:
-    """Apples row-wise softmax to the non-zero elements of the sparse matrix.
+    """Applies row-wise softmax to the non-zero elements of the sparse matrix.
+
+    Equivalently, applies softmax to the non-zero elements of the sparse
+    matrix along the column (``dim=1``) dimension.
 
     If :attr:`input.val` takes shape :attr:`(nnz, D)`, then the output matrix
     :attr:`output` and :attr:`output.val` take the same shape as :attr:`input`