prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN ISSUE ---
  [Sparse] Create a mock implementation in mock_sparse for BSDDMM.
  ## 🔨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/mock_sparse/sddmm.py`
  Content:
  ```
  1 """Sampled Dense-Dense Matrix Multiplication (SDDMM) operator module."""
  2 import torch
  3 
  4 from .sp_matrix import SparseMatrix
  5 
  6 __all__ = ["sddmm"]
  7 
  8 
  9 def sddmm(
  10     A: SparseMatrix, mat1: torch.Tensor, mat2: torch.Tensor
  11 ) -> SparseMatrix:
  12     r"""Sampled-Dense-Dense Matrix Multiplication (SDDMM).
  13 
  14     ``sddmm`` multiplies two dense matrices :attr:``mat1`` and :attr:``mat2``
  15     at the nonzero locations of sparse matrix :attr:``A``. Values of :attr:``A``
  16     is added to the resulting matrix.
  17 
  18     Mathematically ``sddmm`` is formulated as:
  19 
  20     .. math::
  21         out = (mat1 @ mat2) * spy(A) + A
  22 
  23     Parameters
  24     ----------
  25     A : SparseMatrix
  26         Sparse matrix of shape `(M, N)`.
  27     mat1 : Tensor
  28         Dense matrix of shape `(M, K)`
  29     mat2 : Tensor
  30         Dense matrix of shape `(K, N)`
  31 
  32     Returns
  33     -------
  34     SparseMatrix
  35         Sparse matrix of shape `(M, N)`.
  36 
  37     Examples
  38     --------
  39 
  40     >>> row = torch.Tensor([1, 1, 2])
  41     >>> col = torch.Tensor([2, 3, 3])
  42     >>> val = torch.arange(1, 4).float()
  43     >>> A = SparseMatrix(row, col, val, (3, 4))
  44     >>> mat1 = torch.randn(3, 5)
  45     >>> mat2 = torch.randn(5, 4)
  46     >>> dgl.mock_sparse.sddmm(A, mat1, mat2)
  47     SparseMatrix(indices=tensor([[1, 1, 2],
  48             [2, 3, 3]]),
  49     values=tensor([1.8035, 2.3375, 3.1255]),
  50     shape=(3, 4), nnz=3)
  51     """
  52     assert A.val.dim() == 1, (
  53         f"Nonzero elements have values of shape ({A.val.shape[1]}). Expects "
  54         "scalar values. "
  55     )
  56     # PyTorch's sddmm operator only supports CSR format.
  57     res = torch.sparse.sampled_addmm(A.adj.to_sparse_csr(), mat1, mat2)
  58     return SparseMatrix(A.row, A.col, res.values(), A.adj.shape)
  59
  ```
  --- 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/mock_sparse/sddmm.py b/python/dgl/mock_sparse/sddmm.py
  --- a/python/dgl/mock_sparse/sddmm.py
  +++ b/python/dgl/mock_sparse/sddmm.py
  @@ -1,9 +1,9 @@
   """Sampled Dense-Dense Matrix Multiplication (SDDMM) operator module."""
   import torch
   
  -from .sp_matrix import SparseMatrix
  +from .sp_matrix import create_from_coo, SparseMatrix
   
  -__all__ = ["sddmm"]
  +__all__ = ["sddmm", "mock_bsddmm"]
   
   
   def sddmm(
  @@ -56,3 +56,56 @@
       # PyTorch's sddmm operator only supports CSR format.
       res = torch.sparse.sampled_addmm(A.adj.to_sparse_csr(), mat1, mat2)
       return SparseMatrix(A.row, A.col, res.values(), A.adj.shape)
  +
  +
  +def mock_bsddmm(
  +    A: SparseMatrix, mat1: torch.Tensor, mat2: torch.Tensor
  +) -> SparseMatrix:
  +    r"""Batched Sampled-Dense-Dense Matrix Multiplication (SDDMM).
  +
  +    ``bsddmm`` conducts `sddmm` for each batch of the two dense matrices
  +    independently.
  +
  +    In particular, :attr:``mat1`` and :attr:``mat2`` can be 2-D, which will be
  +    reshape as `(B, M, 1)` and `(B, 1, K)` in the computation.
  +
  +    Parameters
  +    ----------
  +    A : SparseMatrix
  +        Sparse matrix of shape `(M, N)`.
  +    mat1 : Tensor
  +        Dense matrix of shape `(B, M, K)` or `(B, M,)`
  +    mat2 : Tensor
  +        Dense matrix of shape `(B, K, N)` or `(B, K,)`
  +
  +    Returns
  +    -------
  +    SparseMatrix
  +        Sparse matrix of shape `(M, N)` with non-zero values of `B` dimension.
  +
  +    Examples
  +    --------
  +
  +    >>> row = torch.tensor([1, 1, 2])
  +    >>> col = torch.tensor([2, 3, 3])
  +    >>> val = torch.arange(1, 4).float()
  +    >>> A = create_from_coo(row, col, val, (3, 4))
  +    >>> mat1 = torch.randn(2, 3, 5)
  +    >>> mat2 = torch.randn(2, 5, 4)
  +    >>> dgl.mock_sparse.mock_bsddmm(A, mat1, mat2)
  +    SparseMatrix(indices=tensor([[1, 1, 2],
  +            [2, 3, 3]]),
  +    values=tensor([[-0.6765, -0.4017],
  +            [ 3.3290,  6.9016],
  +            [ 4.8184,  5.8882]]),
  +    shape=(3, 4), nnz=3)
  +    """
  +    batch_mat1 = [mat1[i, ...] for i in range(mat1.shape[0])]
  +    batch_mat2 = [mat2[i, ...] for i in range(mat2.shape[0])]
  +    batch_ret = [sddmm(A, lhs, rhs) for lhs, rhs in zip(batch_mat1, batch_mat2)]
  +    return create_from_coo(
  +        row=A.row,
  +        col=A.col,
  +        val=torch.stack([sp_mat.val for sp_mat in batch_ret], dim=-1),
  +        shape=A.shape,
  +    )
  
verification_info: {"golden_diff": "diff --git a/python/dgl/mock_sparse/sddmm.py b/python/dgl/mock_sparse/sddmm.py\n--- a/python/dgl/mock_sparse/sddmm.py\n+++ b/python/dgl/mock_sparse/sddmm.py\n@@ -1,9 +1,9 @@\n \"\"\"Sampled Dense-Dense Matrix Multiplication (SDDMM) operator module.\"\"\"\n import torch\n \n-from .sp_matrix import SparseMatrix\n+from .sp_matrix import create_from_coo, SparseMatrix\n \n-__all__ = [\"sddmm\"]\n+__all__ = [\"sddmm\", \"mock_bsddmm\"]\n \n \n def sddmm(\n@@ -56,3 +56,56 @@\n     # PyTorch's sddmm operator only supports CSR format.\n     res = torch.sparse.sampled_addmm(A.adj.to_sparse_csr(), mat1, mat2)\n     return SparseMatrix(A.row, A.col, res.values(), A.adj.shape)\n+\n+\n+def mock_bsddmm(\n+    A: SparseMatrix, mat1: torch.Tensor, mat2: torch.Tensor\n+) -> SparseMatrix:\n+    r\"\"\"Batched Sampled-Dense-Dense Matrix Multiplication (SDDMM).\n+\n+    ``bsddmm`` conducts `sddmm` for each batch of the two dense matrices\n+    independently.\n+\n+    In particular, :attr:``mat1`` and :attr:``mat2`` can be 2-D, which will be\n+    reshape as `(B, M, 1)` and `(B, 1, K)` in the computation.\n+\n+    Parameters\n+    ----------\n+    A : SparseMatrix\n+        Sparse matrix of shape `(M, N)`.\n+    mat1 : Tensor\n+        Dense matrix of shape `(B, M, K)` or `(B, M,)`\n+    mat2 : Tensor\n+        Dense matrix of shape `(B, K, N)` or `(B, K,)`\n+\n+    Returns\n+    -------\n+    SparseMatrix\n+        Sparse matrix of shape `(M, N)` with non-zero values of `B` dimension.\n+\n+    Examples\n+    --------\n+\n+    >>> row = torch.tensor([1, 1, 2])\n+    >>> col = torch.tensor([2, 3, 3])\n+    >>> val = torch.arange(1, 4).float()\n+    >>> A = create_from_coo(row, col, val, (3, 4))\n+    >>> mat1 = torch.randn(2, 3, 5)\n+    >>> mat2 = torch.randn(2, 5, 4)\n+    >>> dgl.mock_sparse.mock_bsddmm(A, mat1, mat2)\n+    SparseMatrix(indices=tensor([[1, 1, 2],\n+            [2, 3, 3]]),\n+    values=tensor([[-0.6765, -0.4017],\n+            [ 3.3290,  6.9016],\n+            [ 4.8184,  5.8882]]),\n+    shape=(3, 4), nnz=3)\n+    \"\"\"\n+    batch_mat1 = [mat1[i, ...] for i in range(mat1.shape[0])]\n+    batch_mat2 = [mat2[i, ...] for i in range(mat2.shape[0])]\n+    batch_ret = [sddmm(A, lhs, rhs) for lhs, rhs in zip(batch_mat1, batch_mat2)]\n+    return create_from_coo(\n+        row=A.row,\n+        col=A.col,\n+        val=torch.stack([sp_mat.val for sp_mat in batch_ret], dim=-1),\n+        shape=A.shape,\n+    )\n", "issue": "[Sparse] Create a mock implementation in mock_sparse for BSDDMM.\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": "\"\"\"Sampled Dense-Dense Matrix Multiplication (SDDMM) operator module.\"\"\"\nimport torch\n\nfrom .sp_matrix import SparseMatrix\n\n__all__ = [\"sddmm\"]\n\n\ndef sddmm(\n    A: SparseMatrix, mat1: torch.Tensor, mat2: torch.Tensor\n) -> SparseMatrix:\n    r\"\"\"Sampled-Dense-Dense Matrix Multiplication (SDDMM).\n\n    ``sddmm`` multiplies two dense matrices :attr:``mat1`` and :attr:``mat2``\n    at the nonzero locations of sparse matrix :attr:``A``. Values of :attr:``A``\n    is added to the resulting matrix.\n\n    Mathematically ``sddmm`` is formulated as:\n\n    .. math::\n        out = (mat1 @ mat2) * spy(A) + A\n\n    Parameters\n    ----------\n    A : SparseMatrix\n        Sparse matrix of shape `(M, N)`.\n    mat1 : Tensor\n        Dense matrix of shape `(M, K)`\n    mat2 : Tensor\n        Dense matrix of shape `(K, N)`\n\n    Returns\n    -------\n    SparseMatrix\n        Sparse matrix of shape `(M, N)`.\n\n    Examples\n    --------\n\n    >>> row = torch.Tensor([1, 1, 2])\n    >>> col = torch.Tensor([2, 3, 3])\n    >>> val = torch.arange(1, 4).float()\n    >>> A = SparseMatrix(row, col, val, (3, 4))\n    >>> mat1 = torch.randn(3, 5)\n    >>> mat2 = torch.randn(5, 4)\n    >>> dgl.mock_sparse.sddmm(A, mat1, mat2)\n    SparseMatrix(indices=tensor([[1, 1, 2],\n            [2, 3, 3]]),\n    values=tensor([1.8035, 2.3375, 3.1255]),\n    shape=(3, 4), nnz=3)\n    \"\"\"\n    assert A.val.dim() == 1, (\n        f\"Nonzero elements have values of shape ({A.val.shape[1]}). Expects \"\n        \"scalar values. \"\n    )\n    # PyTorch's sddmm operator only supports CSR format.\n    res = torch.sparse.sampled_addmm(A.adj.to_sparse_csr(), mat1, mat2)\n    return SparseMatrix(A.row, A.col, res.values(), A.adj.shape)\n", "path": "python/dgl/mock_sparse/sddmm.py"}], "after_files": [{"content": "\"\"\"Sampled Dense-Dense Matrix Multiplication (SDDMM) operator module.\"\"\"\nimport torch\n\nfrom .sp_matrix import create_from_coo, SparseMatrix\n\n__all__ = [\"sddmm\", \"mock_bsddmm\"]\n\n\ndef sddmm(\n    A: SparseMatrix, mat1: torch.Tensor, mat2: torch.Tensor\n) -> SparseMatrix:\n    r\"\"\"Sampled-Dense-Dense Matrix Multiplication (SDDMM).\n\n    ``sddmm`` multiplies two dense matrices :attr:``mat1`` and :attr:``mat2``\n    at the nonzero locations of sparse matrix :attr:``A``. Values of :attr:``A``\n    is added to the resulting matrix.\n\n    Mathematically ``sddmm`` is formulated as:\n\n    .. math::\n        out = (mat1 @ mat2) * spy(A) + A\n\n    Parameters\n    ----------\n    A : SparseMatrix\n        Sparse matrix of shape `(M, N)`.\n    mat1 : Tensor\n        Dense matrix of shape `(M, K)`\n    mat2 : Tensor\n        Dense matrix of shape `(K, N)`\n\n    Returns\n    -------\n    SparseMatrix\n        Sparse matrix of shape `(M, N)`.\n\n    Examples\n    --------\n\n    >>> row = torch.Tensor([1, 1, 2])\n    >>> col = torch.Tensor([2, 3, 3])\n    >>> val = torch.arange(1, 4).float()\n    >>> A = SparseMatrix(row, col, val, (3, 4))\n    >>> mat1 = torch.randn(3, 5)\n    >>> mat2 = torch.randn(5, 4)\n    >>> dgl.mock_sparse.sddmm(A, mat1, mat2)\n    SparseMatrix(indices=tensor([[1, 1, 2],\n            [2, 3, 3]]),\n    values=tensor([1.8035, 2.3375, 3.1255]),\n    shape=(3, 4), nnz=3)\n    \"\"\"\n    assert A.val.dim() == 1, (\n        f\"Nonzero elements have values of shape ({A.val.shape[1]}). Expects \"\n        \"scalar values. \"\n    )\n    # PyTorch's sddmm operator only supports CSR format.\n    res = torch.sparse.sampled_addmm(A.adj.to_sparse_csr(), mat1, mat2)\n    return SparseMatrix(A.row, A.col, res.values(), A.adj.shape)\n\n\ndef mock_bsddmm(\n    A: SparseMatrix, mat1: torch.Tensor, mat2: torch.Tensor\n) -> SparseMatrix:\n    r\"\"\"Batched Sampled-Dense-Dense Matrix Multiplication (SDDMM).\n\n    ``bsddmm`` conducts `sddmm` for each batch of the two dense matrices\n    independently.\n\n    In particular, :attr:``mat1`` and :attr:``mat2`` can be 2-D, which will be\n    reshape as `(B, M, 1)` and `(B, 1, K)` in the computation.\n\n    Parameters\n    ----------\n    A : SparseMatrix\n        Sparse matrix of shape `(M, N)`.\n    mat1 : Tensor\n        Dense matrix of shape `(B, M, K)` or `(B, M,)`\n    mat2 : Tensor\n        Dense matrix of shape `(B, K, N)` or `(B, K,)`\n\n    Returns\n    -------\n    SparseMatrix\n        Sparse matrix of shape `(M, N)` with non-zero values of `B` dimension.\n\n    Examples\n    --------\n\n    >>> row = torch.tensor([1, 1, 2])\n    >>> col = torch.tensor([2, 3, 3])\n    >>> val = torch.arange(1, 4).float()\n    >>> A = create_from_coo(row, col, val, (3, 4))\n    >>> mat1 = torch.randn(2, 3, 5)\n    >>> mat2 = torch.randn(2, 5, 4)\n    >>> dgl.mock_sparse.mock_bsddmm(A, mat1, mat2)\n    SparseMatrix(indices=tensor([[1, 1, 2],\n            [2, 3, 3]]),\n    values=tensor([[-0.6765, -0.4017],\n            [ 3.3290,  6.9016],\n            [ 4.8184,  5.8882]]),\n    shape=(3, 4), nnz=3)\n    \"\"\"\n    batch_mat1 = [mat1[i, ...] for i in range(mat1.shape[0])]\n    batch_mat2 = [mat2[i, ...] for i in range(mat2.shape[0])]\n    batch_ret = [sddmm(A, lhs, rhs) for lhs, rhs in zip(batch_mat1, batch_mat2)]\n    return create_from_coo(\n        row=A.row,\n        col=A.col,\n        val=torch.stack([sp_mat.val for sp_mat in batch_ret], dim=-1),\n        shape=A.shape,\n    )\n", "path": "python/dgl/mock_sparse/sddmm.py"}]}