prompt: 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()
  
  ```
  
golden_diff: 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
  
verification_info: {"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"}]}