Repository: dmlc/dgl

Issue:
[Sparse] Support column-wise softmax.
## 🔨Work Item

**IMPORTANT:**
* This template is only for dev team to track project progress. For feature request or bug report, please use the corresponding issue templates.
* DO NOT create a new work item if the purpose is to fix an existing issue or feature request. We will directly use the issue in the project tracker.

Project tracker: https://github.com/orgs/dmlc/projects/2

## Description

<!-- short description of the work item -->

## Depending work items or issues

<!-- what must be done before this -->



Discussion:

Resolving patch:
diff --git a/python/dgl/sparse/softmax.py b/python/dgl/sparse/softmax.py
--- a/python/dgl/sparse/softmax.py
+++ b/python/dgl/sparse/softmax.py
@@ -8,11 +8,10 @@
 __all__ = ["softmax"]
 
 
-def softmax(input: SparseMatrix) -> SparseMatrix:
-    """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.
+def softmax(input: SparseMatrix, dim: int = 1) -> SparseMatrix:
+    """Applies softmax to the non-zero elements of the sparse matrix on the
+    dimension :attr:``dim``. dim = 0 or 1 indicates column-wise or row-wise
+    softmax respectively.
 
     If :attr:`input.val` takes shape ``(nnz, D)``, then the output matrix
     :attr:`output` and :attr:`output.val` take the same shape as :attr:`input`
@@ -32,11 +31,10 @@ def softmax(input: SparseMatrix) -> SparseMatrix:
     Examples
     --------
 
-    Case1: matrix with values of shape (nnz)
+    Case1: row-wise softmax on matrix with values of shape (nnz)
 
     >>> indices = torch.tensor([[0, 0, 1, 2], [1, 2, 2, 0]])
-    >>> nnz = len(row)
-    >>> val = torch.arange(nnz).float()
+    >>> val = torch.tensor([0., 1., 2., 3.])
     >>> A = dglsp.spmatrix(indices, val)
     >>> dglsp.softmax(A)
     SparseMatrix(indices=tensor([[0, 0, 1, 2],
@@ -44,7 +42,7 @@ def softmax(input: SparseMatrix) -> SparseMatrix:
                  values=tensor([0.2689, 0.7311, 1.0000, 1.0000]),
                  shape=(3, 3), nnz=4)
 
-    Case2: matrix with values of shape (nnz, D)
+    Case2: row-wise softmax on matrix with values of shape (nnz, D)
 
     >>> indices = torch.tensor([[0, 0, 1, 2], [1, 2, 2, 0]])
     >>> val = torch.tensor([[0., 7.], [1., 3.], [2., 2.], [3., 1.]])
@@ -57,8 +55,21 @@ def softmax(input: SparseMatrix) -> SparseMatrix:
                                 [1.0000, 1.0000],
                                 [1.0000, 1.0000]]),
                  shape=(3, 3), nnz=4, val_size=(2,))
+
+    Case3: column-wise softmax on matrix with values of shape (nnz)
+
+    >>> indices = torch.tensor([[0, 0, 1, 2], [1, 2, 2, 0]])
+    >>> val = torch.tensor([0., 1., 2., 3.])
+    >>> A = dglsp.spmatrix(indices, val)
+    >>> dglsp.softmax(A, 0)
+    SparseMatrix(indices=tensor([[0, 0, 1, 2],
+                                 [1, 2, 2, 0]]),
+                 values=tensor([1.0000, 0.2689, 0.7311, 1.0000]),
+                 shape=(3, 3), nnz=4)
     """
-    return SparseMatrix(torch.ops.dgl_sparse.softmax(input.c_sparse_matrix))
+    return SparseMatrix(
+        torch.ops.dgl_sparse.softmax(input.c_sparse_matrix, dim)
+    )
 
 
 SparseMatrix.softmax = softmax


Test patch:
diff --git a/tests/python/pytorch/sparse/test_softmax.py b/tests/python/pytorch/sparse/test_softmax.py
--- a/tests/python/pytorch/sparse/test_softmax.py
+++ b/tests/python/pytorch/sparse/test_softmax.py
@@ -10,7 +10,8 @@
 
 @pytest.mark.parametrize("val_D", [None, 2])
 @pytest.mark.parametrize("csr", [True, False])
-def test_softmax(val_D, csr):
+@pytest.mark.parametrize("dim", [0, 1])
+def test_softmax(val_D, csr, dim):
     dev = F.ctx()
     row = torch.tensor([0, 0, 1, 1]).to(dev)
     col = torch.tensor([0, 2, 1, 2]).to(dev)
@@ -27,8 +28,11 @@ def test_softmax(val_D, csr):
         # Test CSR
         A.csr()
 
-    A_max = softmax(A)
-    g = dgl.graph((col, row), num_nodes=max(A.shape))
+    A_max = softmax(A, dim)
+    if dim == 1:
+        g = dgl.graph((col, row), num_nodes=max(A.shape))
+    else:
+        g = dgl.graph((row, col), num_nodes=max(A.shape))
     val_g = val.clone().requires_grad_()
     score = dgl.nn.functional.edge_softmax(g, val_g)
     assert torch.allclose(A_max.val, score)