We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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 -->


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `python/dgl/sparse/softmax.py`
Content:
```
1 """Softmax op for SparseMatrix"""
2 # pylint: disable=invalid-name, W0622
3 
4 import torch
5 
6 from .sparse_matrix import SparseMatrix
7 
8 __all__ = ["softmax"]
9 
10 
11 def softmax(input: SparseMatrix) -> SparseMatrix:
12     """Applies row-wise softmax to the non-zero elements of the sparse matrix.
13 
14     Equivalently, applies softmax to the non-zero elements of the sparse
15     matrix along the column (``dim=1``) dimension.
16 
17     If :attr:`input.val` takes shape ``(nnz, D)``, then the output matrix
18     :attr:`output` and :attr:`output.val` take the same shape as :attr:`input`
19     and :attr:`input.val`. :attr:`output.val[:, i]` is calculated based on
20     :attr:`input.val[:, i]`.
21 
22     Parameters
23     ----------
24     input : SparseMatrix
25         The input sparse matrix
26 
27     Returns
28     -------
29     SparseMatrix
30         The output sparse matrix
31 
32     Examples
33     --------
34 
35     Case1: matrix with values of shape (nnz)
36 
37     >>> indices = torch.tensor([[0, 0, 1, 2], [1, 2, 2, 0]])
38     >>> nnz = len(row)
39     >>> val = torch.arange(nnz).float()
40     >>> A = dglsp.spmatrix(indices, val)
41     >>> dglsp.softmax(A)
42     SparseMatrix(indices=tensor([[0, 0, 1, 2],
43                                  [1, 2, 2, 0]]),
44                  values=tensor([0.2689, 0.7311, 1.0000, 1.0000]),
45                  shape=(3, 3), nnz=4)
46 
47     Case2: matrix with values of shape (nnz, D)
48 
49     >>> indices = torch.tensor([[0, 0, 1, 2], [1, 2, 2, 0]])
50     >>> val = torch.tensor([[0., 7.], [1., 3.], [2., 2.], [3., 1.]])
51     >>> A = dglsp.spmatrix(indices, val)
52     >>> dglsp.softmax(A)
53     SparseMatrix(indices=tensor([[0, 0, 1, 2],
54                                  [1, 2, 2, 0]]),
55                  values=tensor([[0.2689, 0.9820],
56                                 [0.7311, 0.0180],
57                                 [1.0000, 1.0000],
58                                 [1.0000, 1.0000]]),
59                  shape=(3, 3), nnz=4, val_size=(2,))
60     """
61     return SparseMatrix(torch.ops.dgl_sparse.softmax(input.c_sparse_matrix))
62 
63 
64 SparseMatrix.softmax = softmax
65
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 @@
     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 @@
                  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 @@
                                 [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

{"golden_diff": "diff --git a/python/dgl/sparse/softmax.py b/python/dgl/sparse/softmax.py\n--- a/python/dgl/sparse/softmax.py\n+++ b/python/dgl/sparse/softmax.py\n@@ -8,11 +8,10 @@\n __all__ = [\"softmax\"]\n \n \n-def softmax(input: SparseMatrix) -> SparseMatrix:\n-    \"\"\"Applies row-wise softmax to the non-zero elements of the sparse matrix.\n-\n-    Equivalently, applies softmax to the non-zero elements of the sparse\n-    matrix along the column (``dim=1``) dimension.\n+def softmax(input: SparseMatrix, dim: int = 1) -> SparseMatrix:\n+    \"\"\"Applies softmax to the non-zero elements of the sparse matrix on the\n+    dimension :attr:``dim``. dim = 0 or 1 indicates column-wise or row-wise\n+    softmax respectively.\n \n     If :attr:`input.val` takes shape ``(nnz, D)``, then the output matrix\n     :attr:`output` and :attr:`output.val` take the same shape as :attr:`input`\n@@ -32,11 +31,10 @@\n     Examples\n     --------\n \n-    Case1: matrix with values of shape (nnz)\n+    Case1: row-wise softmax on matrix with values of shape (nnz)\n \n     >>> indices = torch.tensor([[0, 0, 1, 2], [1, 2, 2, 0]])\n-    >>> nnz = len(row)\n-    >>> val = torch.arange(nnz).float()\n+    >>> val = torch.tensor([0., 1., 2., 3.])\n     >>> A = dglsp.spmatrix(indices, val)\n     >>> dglsp.softmax(A)\n     SparseMatrix(indices=tensor([[0, 0, 1, 2],\n@@ -44,7 +42,7 @@\n                  values=tensor([0.2689, 0.7311, 1.0000, 1.0000]),\n                  shape=(3, 3), nnz=4)\n \n-    Case2: matrix with values of shape (nnz, D)\n+    Case2: row-wise softmax on matrix with values of shape (nnz, D)\n \n     >>> indices = torch.tensor([[0, 0, 1, 2], [1, 2, 2, 0]])\n     >>> val = torch.tensor([[0., 7.], [1., 3.], [2., 2.], [3., 1.]])\n@@ -57,8 +55,21 @@\n                                 [1.0000, 1.0000],\n                                 [1.0000, 1.0000]]),\n                  shape=(3, 3), nnz=4, val_size=(2,))\n+\n+    Case3: column-wise softmax on matrix with values of shape (nnz)\n+\n+    >>> indices = torch.tensor([[0, 0, 1, 2], [1, 2, 2, 0]])\n+    >>> val = torch.tensor([0., 1., 2., 3.])\n+    >>> A = dglsp.spmatrix(indices, val)\n+    >>> dglsp.softmax(A, 0)\n+    SparseMatrix(indices=tensor([[0, 0, 1, 2],\n+                                 [1, 2, 2, 0]]),\n+                 values=tensor([1.0000, 0.2689, 0.7311, 1.0000]),\n+                 shape=(3, 3), nnz=4)\n     \"\"\"\n-    return SparseMatrix(torch.ops.dgl_sparse.softmax(input.c_sparse_matrix))\n+    return SparseMatrix(\n+        torch.ops.dgl_sparse.softmax(input.c_sparse_matrix, dim)\n+    )\n \n \n SparseMatrix.softmax = softmax\n", "issue": "[Sparse] Support column-wise softmax.\n## \ud83d\udd28Work Item\r\n\r\n**IMPORTANT:**\r\n* This template is only for dev team to track project progress. For feature request or bug report, please use the corresponding issue templates.\r\n* 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.\r\n\r\nProject tracker: https://github.com/orgs/dmlc/projects/2\r\n\r\n## Description\r\n\r\n<!-- short description of the work item -->\r\n\r\n## Depending work items or issues\r\n\r\n<!-- what must be done before this -->\r\n\n", "before_files": [{"content": "\"\"\"Softmax op for SparseMatrix\"\"\"\n# pylint: disable=invalid-name, W0622\n\nimport torch\n\nfrom .sparse_matrix import SparseMatrix\n\n__all__ = [\"softmax\"]\n\n\ndef softmax(input: SparseMatrix) -> SparseMatrix:\n    \"\"\"Applies row-wise softmax to the non-zero elements of the sparse matrix.\n\n    Equivalently, applies softmax to the non-zero elements of the sparse\n    matrix along the column (``dim=1``) dimension.\n\n    If :attr:`input.val` takes shape ``(nnz, D)``, then the output matrix\n    :attr:`output` and :attr:`output.val` take the same shape as :attr:`input`\n    and :attr:`input.val`. :attr:`output.val[:, i]` is calculated based on\n    :attr:`input.val[:, i]`.\n\n    Parameters\n    ----------\n    input : SparseMatrix\n        The input sparse matrix\n\n    Returns\n    -------\n    SparseMatrix\n        The output sparse matrix\n\n    Examples\n    --------\n\n    Case1: matrix with values of shape (nnz)\n\n    >>> indices = torch.tensor([[0, 0, 1, 2], [1, 2, 2, 0]])\n    >>> nnz = len(row)\n    >>> val = torch.arange(nnz).float()\n    >>> A = dglsp.spmatrix(indices, val)\n    >>> dglsp.softmax(A)\n    SparseMatrix(indices=tensor([[0, 0, 1, 2],\n                                 [1, 2, 2, 0]]),\n                 values=tensor([0.2689, 0.7311, 1.0000, 1.0000]),\n                 shape=(3, 3), nnz=4)\n\n    Case2: matrix with values of shape (nnz, D)\n\n    >>> indices = torch.tensor([[0, 0, 1, 2], [1, 2, 2, 0]])\n    >>> val = torch.tensor([[0., 7.], [1., 3.], [2., 2.], [3., 1.]])\n    >>> A = dglsp.spmatrix(indices, val)\n    >>> dglsp.softmax(A)\n    SparseMatrix(indices=tensor([[0, 0, 1, 2],\n                                 [1, 2, 2, 0]]),\n                 values=tensor([[0.2689, 0.9820],\n                                [0.7311, 0.0180],\n                                [1.0000, 1.0000],\n                                [1.0000, 1.0000]]),\n                 shape=(3, 3), nnz=4, val_size=(2,))\n    \"\"\"\n    return SparseMatrix(torch.ops.dgl_sparse.softmax(input.c_sparse_matrix))\n\n\nSparseMatrix.softmax = softmax\n", "path": "python/dgl/sparse/softmax.py"}], "after_files": [{"content": "\"\"\"Softmax op for SparseMatrix\"\"\"\n# pylint: disable=invalid-name, W0622\n\nimport torch\n\nfrom .sparse_matrix import SparseMatrix\n\n__all__ = [\"softmax\"]\n\n\ndef softmax(input: SparseMatrix, dim: int = 1) -> SparseMatrix:\n    \"\"\"Applies softmax to the non-zero elements of the sparse matrix on the\n    dimension :attr:``dim``. dim = 0 or 1 indicates column-wise or row-wise\n    softmax respectively.\n\n    If :attr:`input.val` takes shape ``(nnz, D)``, then the output matrix\n    :attr:`output` and :attr:`output.val` take the same shape as :attr:`input`\n    and :attr:`input.val`. :attr:`output.val[:, i]` is calculated based on\n    :attr:`input.val[:, i]`.\n\n    Parameters\n    ----------\n    input : SparseMatrix\n        The input sparse matrix\n\n    Returns\n    -------\n    SparseMatrix\n        The output sparse matrix\n\n    Examples\n    --------\n\n    Case1: row-wise softmax on matrix with values of shape (nnz)\n\n    >>> indices = torch.tensor([[0, 0, 1, 2], [1, 2, 2, 0]])\n    >>> val = torch.tensor([0., 1., 2., 3.])\n    >>> A = dglsp.spmatrix(indices, val)\n    >>> dglsp.softmax(A)\n    SparseMatrix(indices=tensor([[0, 0, 1, 2],\n                                 [1, 2, 2, 0]]),\n                 values=tensor([0.2689, 0.7311, 1.0000, 1.0000]),\n                 shape=(3, 3), nnz=4)\n\n    Case2: row-wise softmax on matrix with values of shape (nnz, D)\n\n    >>> indices = torch.tensor([[0, 0, 1, 2], [1, 2, 2, 0]])\n    >>> val = torch.tensor([[0., 7.], [1., 3.], [2., 2.], [3., 1.]])\n    >>> A = dglsp.spmatrix(indices, val)\n    >>> dglsp.softmax(A)\n    SparseMatrix(indices=tensor([[0, 0, 1, 2],\n                                 [1, 2, 2, 0]]),\n                 values=tensor([[0.2689, 0.9820],\n                                [0.7311, 0.0180],\n                                [1.0000, 1.0000],\n                                [1.0000, 1.0000]]),\n                 shape=(3, 3), nnz=4, val_size=(2,))\n\n    Case3: column-wise softmax on matrix with values of shape (nnz)\n\n    >>> indices = torch.tensor([[0, 0, 1, 2], [1, 2, 2, 0]])\n    >>> val = torch.tensor([0., 1., 2., 3.])\n    >>> A = dglsp.spmatrix(indices, val)\n    >>> dglsp.softmax(A, 0)\n    SparseMatrix(indices=tensor([[0, 0, 1, 2],\n                                 [1, 2, 2, 0]]),\n                 values=tensor([1.0000, 0.2689, 0.7311, 1.0000]),\n                 shape=(3, 3), nnz=4)\n    \"\"\"\n    return SparseMatrix(\n        torch.ops.dgl_sparse.softmax(input.c_sparse_matrix, dim)\n    )\n\n\nSparseMatrix.softmax = softmax\n", "path": "python/dgl/sparse/softmax.py"}]}