Repository: dmlc/dgl

Issue:
[Sparse] Support 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/mock_sparse2/__init__.py b/python/dgl/mock_sparse2/__init__.py
--- a/python/dgl/mock_sparse2/__init__.py
+++ b/python/dgl/mock_sparse2/__init__.py
@@ -12,6 +12,7 @@
 from .matmul import *
 from .reduction import *  # pylint: disable=W0622
 from .sddmm import *
+from .softmax import *
 from .sparse_matrix import *
 from .unary_op_diag import *
 from .unary_op_sp import *
diff --git a/python/dgl/mock_sparse2/softmax.py b/python/dgl/mock_sparse2/softmax.py
new file mode 100644
--- /dev/null
+++ b/python/dgl/mock_sparse2/softmax.py
@@ -0,0 +1,58 @@
+"""Softmax op for SparseMatrix"""
+# pylint: disable=invalid-name
+
+import torch
+
+from .sparse_matrix import SparseMatrix
+
+__all__ = ["softmax"]
+
+
+def softmax(A: SparseMatrix) -> SparseMatrix:
+    """Apply row-wise softmax to the non-zero entries of the sparse matrix.
+
+    If :attr:`A.val` takes shape :attr:`(nnz, D)`, then the output matrix
+    :attr:`A'` and :attr:`A'.val` take the same shape as :attr:`A` and
+    :attr:`A.val`. :attr:`A'.val[:, i]` is calculated based on
+    :attr:`A.val[:, i]`.
+
+    Parameters
+    ----------
+    A : SparseMatrix
+        The input sparse matrix
+
+    Returns
+    -------
+    SparseMatrix
+        The output sparse matrix
+
+    Examples
+    --------
+
+    Case1: matrix with values of shape (nnz)
+
+    >>> row = torch.tensor([0, 0, 1, 2])
+    >>> col = torch.tensor([1, 2, 2, 0])
+    >>> nnz = len(row)
+    >>> val = torch.arange(nnz).float()
+    >>> A = create_from_coo(row, col, val)
+    >>> softmax(A)
+    SparseMatrix(indices=tensor([[0, 0, 1, 2],
+        [1, 2, 2, 0]]),
+    values=tensor([0.2689, 0.7311, 1.0000, 1.0000]),
+    shape=(3, 3), nnz=4)
+
+    Case2: matrix with values of shape (nnz, D)
+
+    >>> val = torch.tensor([[0., 7.], [1., 3.], [2., 2.], [3., 1.]])
+    >>> A = create_from_coo(row, col, val)
+    >>> softmax(A)
+    SparseMatrix(indices=tensor([[0, 0, 1, 2],
+        [1, 2, 2, 0]]),
+    values=tensor([[0.2689, 0.9820],
+        [0.7311, 0.0180],
+        [1.0000, 1.0000],
+        [1.0000, 1.0000]]),
+    shape=(3, 3), nnz=4)
+    """
+    return SparseMatrix(torch.ops.dgl_sparse.softmax(A.c_sparse_matrix))


Test patch:
diff --git a/tests/pytorch/mock_sparse2/test_softmax.py b/tests/pytorch/mock_sparse2/test_softmax.py
new file mode 100644
--- /dev/null
+++ b/tests/pytorch/mock_sparse2/test_softmax.py
@@ -0,0 +1,44 @@
+import sys
+
+import backend as F
+
+import dgl
+import pytest
+import torch
+
+from dgl.mock_sparse2 import create_from_coo, softmax
+
+# TODO(#4818): Skipping tests on win.
+if not sys.platform.startswith("linux"):
+    pytest.skip("skipping tests on win", allow_module_level=True)
+
+
+@pytest.mark.parametrize("val_D", [None, 2])
+@pytest.mark.parametrize("csr", [True, False])
+def test_softmax(val_D, csr):
+    dev = F.ctx()
+    row = torch.tensor([0, 0, 1, 1]).to(dev)
+    col = torch.tensor([0, 2, 1, 2]).to(dev)
+    nnz = len(row)
+    if val_D is None:
+        val = torch.randn(nnz).to(dev)
+    else:
+        val = torch.randn(nnz, val_D).to(dev)
+
+    val_sparse = val.clone().requires_grad_()
+    A = create_from_coo(row, col, val_sparse)
+
+    if csr:
+        # Test CSR
+        A.csr()
+
+    A_max = softmax(A)
+    g = dgl.graph((col, row), 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)
+
+    grad = torch.randn_like(score).to(dev)
+    A_max.val.backward(grad)
+    score.backward(grad)
+    assert torch.allclose(A.val.grad, val_g.grad)