We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[Sparse] Support SparseMatrix element-wise multiplication with different sparsities.
## 🔨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

<del>
Breakdown into two work items.

1. Support on CPU (#5519).

3. Support on GPU.

</del>

Implemented by concatenation and unique on two COO matrices.
## 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/elementwise_op_sp.py`
Content:
```
1 """DGL elementwise operators for sparse matrix module."""
2 from typing import Union
3 
4 import torch
5 
6 from .sparse_matrix import diag, SparseMatrix, val_like
7 from .utils import is_scalar, Scalar
8 
9 
10 def spsp_add(A, B):
11     """Invoke C++ sparse library for addition"""
12     return SparseMatrix(
13         torch.ops.dgl_sparse.spsp_add(A.c_sparse_matrix, B.c_sparse_matrix)
14     )
15 
16 
17 def sp_add(A: SparseMatrix, B: SparseMatrix) -> SparseMatrix:
18     """Elementwise addition
19 
20     Parameters
21     ----------
22     A : SparseMatrix
23         Sparse matrix
24     B : SparseMatrix
25         Sparse matrix
26 
27     Returns
28     -------
29     SparseMatrix
30         Sparse matrix
31 
32     Examples
33     --------
34 
35     >>> indices = torch.tensor([[1, 0, 2], [0, 3, 2]])
36     >>> val = torch.tensor([10, 20, 30])
37     >>> A = dglsp.spmatrix(indices, val, shape=(3, 4))
38     >>> A + A
39     SparseMatrix(indices=tensor([[0, 1, 2],
40                                  [3, 0, 2]]),
41                  values=tensor([40, 20, 60]),
42                  shape=(3, 4), nnz=3)
43     """
44     # Python falls back to B.__radd__ then TypeError when NotImplemented is
45     # returned.
46     return spsp_add(A, B) if isinstance(B, SparseMatrix) else NotImplemented
47 
48 
49 def sp_sub(A: SparseMatrix, B: SparseMatrix) -> SparseMatrix:
50     """Elementwise subtraction
51 
52     Parameters
53     ----------
54     A : SparseMatrix
55         Sparse matrix
56     B : SparseMatrix
57         Sparse matrix
58 
59     Returns
60     -------
61     SparseMatrix
62         Sparse matrix
63 
64     Examples
65     --------
66 
67     >>> indices = torch.tensor([[1, 0, 2], [0, 3, 2]])
68     >>> val = torch.tensor([10, 20, 30])
69     >>> val2 = torch.tensor([5, 10, 15])
70     >>> A = dglsp.spmatrix(indices, val, shape=(3, 4))
71     >>> B = dglsp.spmatrix(indices, val2, shape=(3, 4))
72     >>> A - B
73     SparseMatrix(indices=tensor([[0, 1, 2],
74                                  [3, 0, 2]]),
75                  values=tensor([10, 5, 15]),
76                  shape=(3, 4), nnz=3)
77     """
78     # Python falls back to B.__rsub__ then TypeError when NotImplemented is
79     # returned.
80     return spsp_add(A, -B) if isinstance(B, SparseMatrix) else NotImplemented
81 
82 
83 def sp_mul(A: SparseMatrix, B: Union[SparseMatrix, Scalar]) -> SparseMatrix:
84     """Elementwise multiplication
85 
86     If :attr:`B` is a sparse matrix, both :attr:`A` and :attr:`B` must be
87     diagonal matrices.
88 
89     Parameters
90     ----------
91     A : SparseMatrix
92         First operand
93     B : SparseMatrix or Scalar
94         Second operand
95 
96     Returns
97     -------
98     SparseMatrix
99         Result of A * B
100 
101     Examples
102     --------
103 
104     >>> indices = torch.tensor([[1, 0, 2], [0, 3, 2]])
105     >>> val = torch.tensor([1, 2, 3])
106     >>> A = dglsp.spmatrix(indices, val, shape=(3, 4))
107 
108     >>> A * 2
109     SparseMatrix(indices=tensor([[1, 0, 2],
110                                  [0, 3, 2]]),
111     values=tensor([2, 4, 6]),
112     shape=(3, 4), nnz=3)
113 
114     >>> 2 * A
115     SparseMatrix(indices=tensor([[1, 0, 2],
116                                  [0, 3, 2]]),
117     values=tensor([2, 4, 6]),
118     shape=(3, 4), nnz=3)
119     """
120     if is_scalar(B):
121         return val_like(A, A.val * B)
122     if A.is_diag() and B.is_diag():
123         assert A.shape == B.shape, (
124             f"The shape of diagonal matrix A {A.shape} and B {B.shape} must"
125             f"match for elementwise multiplication."
126         )
127         return diag(A.val * B.val, A.shape)
128     # Python falls back to B.__rmul__(A) then TypeError when NotImplemented is
129     # returned.
130     # So this also handles the case of scalar * SparseMatrix since we set
131     # SparseMatrix.__rmul__ to be the same as SparseMatrix.__mul__.
132     return NotImplemented
133 
134 
135 def sp_div(A: SparseMatrix, B: Union[SparseMatrix, Scalar]) -> SparseMatrix:
136     """Elementwise division
137 
138     If :attr:`B` is a sparse matrix, both :attr:`A` and :attr:`B` must be
139     diagonal matrices.
140 
141     Parameters
142     ----------
143     A : SparseMatrix
144         First operand
145     B : SparseMatrix or Scalar
146         Second operand
147 
148     Returns
149     -------
150     SparseMatrix
151         Result of A / B
152 
153     Examples
154     --------
155     >>> indices = torch.tensor([[1, 0, 2], [0, 3, 2]])
156     >>> val = torch.tensor([1, 2, 3])
157     >>> A = dglsp.spmatrix(indices, val, shape=(3, 4))
158     >>> A / 2
159     SparseMatrix(indices=tensor([[1, 0, 2],
160                                  [0, 3, 2]]),
161                  values=tensor([0.5000, 1.0000, 1.5000]),
162                  shape=(3, 4), nnz=3)
163     """
164     if is_scalar(B):
165         return val_like(A, A.val / B)
166     if A.is_diag() and B.is_diag():
167         assert A.shape == B.shape, (
168             f"The shape of diagonal matrix A {A.shape} and B {B.shape} must"
169             f"match for elementwise division."
170         )
171         return diag(A.val / B.val, A.shape)
172     # Python falls back to B.__rtruediv__(A) then TypeError when NotImplemented
173     # is returned.
174     return NotImplemented
175 
176 
177 def sp_power(A: SparseMatrix, scalar: Scalar) -> SparseMatrix:
178     """Take the power of each nonzero element and return a sparse matrix with
179     the result.
180 
181     Parameters
182     ----------
183     A : SparseMatrix
184         Sparse matrix
185     scalar : float or int
186         Exponent
187 
188     Returns
189     -------
190     SparseMatrix
191         Sparse matrix
192 
193     Examples
194     --------
195     >>> indices = torch.tensor([[1, 0, 2], [0, 3, 2]])
196     >>> val = torch.tensor([10, 20, 30])
197     >>> A = dglsp.spmatrix(indices, val)
198     >>> A ** 2
199     SparseMatrix(indices=tensor([[1, 0, 2],
200                                  [0, 3, 2]]),
201     values=tensor([100, 400, 900]),
202     shape=(3, 4), nnz=3)
203     """
204     # Python falls back to scalar.__rpow__ then TypeError when NotImplemented
205     # is returned.
206     return val_like(A, A.val**scalar) if is_scalar(scalar) else NotImplemented
207 
208 
209 SparseMatrix.__add__ = sp_add
210 SparseMatrix.__sub__ = sp_sub
211 SparseMatrix.__mul__ = sp_mul
212 SparseMatrix.__rmul__ = sp_mul
213 SparseMatrix.__truediv__ = sp_div
214 SparseMatrix.__pow__ = sp_power
215
```
--- 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/elementwise_op_sp.py b/python/dgl/sparse/elementwise_op_sp.py
--- a/python/dgl/sparse/elementwise_op_sp.py
+++ b/python/dgl/sparse/elementwise_op_sp.py
@@ -14,6 +14,13 @@
     )
 
 
+def spsp_mul(A, B):
+    """Invoke C++ sparse library for multiplication"""
+    return SparseMatrix(
+        torch.ops.dgl_sparse.spsp_mul(A.c_sparse_matrix, B.c_sparse_matrix)
+    )
+
+
 def sp_add(A: SparseMatrix, B: SparseMatrix) -> SparseMatrix:
     """Elementwise addition
 
@@ -83,8 +90,8 @@
 def sp_mul(A: SparseMatrix, B: Union[SparseMatrix, Scalar]) -> SparseMatrix:
     """Elementwise multiplication
 
-    If :attr:`B` is a sparse matrix, both :attr:`A` and :attr:`B` must be
-    diagonal matrices.
+    Note that if both :attr:`A` and :attr:`B` are sparse matrices, both of them
+    need to be diagonal or on CPU.
 
     Parameters
     ----------
@@ -108,28 +115,27 @@
     >>> A * 2
     SparseMatrix(indices=tensor([[1, 0, 2],
                                  [0, 3, 2]]),
-    values=tensor([2, 4, 6]),
-    shape=(3, 4), nnz=3)
+                 values=tensor([2, 4, 6]),
+                 shape=(3, 4), nnz=3)
 
     >>> 2 * A
     SparseMatrix(indices=tensor([[1, 0, 2],
                                  [0, 3, 2]]),
-    values=tensor([2, 4, 6]),
-    shape=(3, 4), nnz=3)
+                 values=tensor([2, 4, 6]),
+                 shape=(3, 4), nnz=3)
+
+    >>> indices2 = torch.tensor([[2, 0, 1], [0, 3, 2]])
+    >>> val2 = torch.tensor([3, 2, 1])
+    >>> B = dglsp.spmatrix(indices2, val2, shape=(3, 4))
+    >>> A * B
+    SparseMatrix(indices=tensor([[0],
+                                 [3]]),
+                 values=tensor([4]),
+                 shape=(3, 4), nnz=1)
     """
     if is_scalar(B):
         return val_like(A, A.val * B)
-    if A.is_diag() and B.is_diag():
-        assert A.shape == B.shape, (
-            f"The shape of diagonal matrix A {A.shape} and B {B.shape} must"
-            f"match for elementwise multiplication."
-        )
-        return diag(A.val * B.val, A.shape)
-    # Python falls back to B.__rmul__(A) then TypeError when NotImplemented is
-    # returned.
-    # So this also handles the case of scalar * SparseMatrix since we set
-    # SparseMatrix.__rmul__ to be the same as SparseMatrix.__mul__.
-    return NotImplemented
+    return spsp_mul(A, B)
 
 
 def sp_div(A: SparseMatrix, B: Union[SparseMatrix, Scalar]) -> SparseMatrix:

{"golden_diff": "diff --git a/python/dgl/sparse/elementwise_op_sp.py b/python/dgl/sparse/elementwise_op_sp.py\n--- a/python/dgl/sparse/elementwise_op_sp.py\n+++ b/python/dgl/sparse/elementwise_op_sp.py\n@@ -14,6 +14,13 @@\n     )\n \n \n+def spsp_mul(A, B):\n+    \"\"\"Invoke C++ sparse library for multiplication\"\"\"\n+    return SparseMatrix(\n+        torch.ops.dgl_sparse.spsp_mul(A.c_sparse_matrix, B.c_sparse_matrix)\n+    )\n+\n+\n def sp_add(A: SparseMatrix, B: SparseMatrix) -> SparseMatrix:\n     \"\"\"Elementwise addition\n \n@@ -83,8 +90,8 @@\n def sp_mul(A: SparseMatrix, B: Union[SparseMatrix, Scalar]) -> SparseMatrix:\n     \"\"\"Elementwise multiplication\n \n-    If :attr:`B` is a sparse matrix, both :attr:`A` and :attr:`B` must be\n-    diagonal matrices.\n+    Note that if both :attr:`A` and :attr:`B` are sparse matrices, both of them\n+    need to be diagonal or on CPU.\n \n     Parameters\n     ----------\n@@ -108,28 +115,27 @@\n     >>> A * 2\n     SparseMatrix(indices=tensor([[1, 0, 2],\n                                  [0, 3, 2]]),\n-    values=tensor([2, 4, 6]),\n-    shape=(3, 4), nnz=3)\n+                 values=tensor([2, 4, 6]),\n+                 shape=(3, 4), nnz=3)\n \n     >>> 2 * A\n     SparseMatrix(indices=tensor([[1, 0, 2],\n                                  [0, 3, 2]]),\n-    values=tensor([2, 4, 6]),\n-    shape=(3, 4), nnz=3)\n+                 values=tensor([2, 4, 6]),\n+                 shape=(3, 4), nnz=3)\n+\n+    >>> indices2 = torch.tensor([[2, 0, 1], [0, 3, 2]])\n+    >>> val2 = torch.tensor([3, 2, 1])\n+    >>> B = dglsp.spmatrix(indices2, val2, shape=(3, 4))\n+    >>> A * B\n+    SparseMatrix(indices=tensor([[0],\n+                                 [3]]),\n+                 values=tensor([4]),\n+                 shape=(3, 4), nnz=1)\n     \"\"\"\n     if is_scalar(B):\n         return val_like(A, A.val * B)\n-    if A.is_diag() and B.is_diag():\n-        assert A.shape == B.shape, (\n-            f\"The shape of diagonal matrix A {A.shape} and B {B.shape} must\"\n-            f\"match for elementwise multiplication.\"\n-        )\n-        return diag(A.val * B.val, A.shape)\n-    # Python falls back to B.__rmul__(A) then TypeError when NotImplemented is\n-    # returned.\n-    # So this also handles the case of scalar * SparseMatrix since we set\n-    # SparseMatrix.__rmul__ to be the same as SparseMatrix.__mul__.\n-    return NotImplemented\n+    return spsp_mul(A, B)\n \n \n def sp_div(A: SparseMatrix, B: Union[SparseMatrix, Scalar]) -> SparseMatrix:\n", "issue": "[Sparse] Support SparseMatrix element-wise multiplication with different sparsities.\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<del>\r\nBreakdown into two work items.\r\n\r\n1. Support on CPU (#5519).\r\n\r\n3. Support on GPU.\r\n\r\n</del>\r\n\r\nImplemented by concatenation and unique on two COO matrices.\r\n## Depending work items or issues\r\n\r\n<!-- what must be done before this -->\r\n\n", "before_files": [{"content": "\"\"\"DGL elementwise operators for sparse matrix module.\"\"\"\nfrom typing import Union\n\nimport torch\n\nfrom .sparse_matrix import diag, SparseMatrix, val_like\nfrom .utils import is_scalar, Scalar\n\n\ndef spsp_add(A, B):\n    \"\"\"Invoke C++ sparse library for addition\"\"\"\n    return SparseMatrix(\n        torch.ops.dgl_sparse.spsp_add(A.c_sparse_matrix, B.c_sparse_matrix)\n    )\n\n\ndef sp_add(A: SparseMatrix, B: SparseMatrix) -> SparseMatrix:\n    \"\"\"Elementwise addition\n\n    Parameters\n    ----------\n    A : SparseMatrix\n        Sparse matrix\n    B : SparseMatrix\n        Sparse matrix\n\n    Returns\n    -------\n    SparseMatrix\n        Sparse matrix\n\n    Examples\n    --------\n\n    >>> indices = torch.tensor([[1, 0, 2], [0, 3, 2]])\n    >>> val = torch.tensor([10, 20, 30])\n    >>> A = dglsp.spmatrix(indices, val, shape=(3, 4))\n    >>> A + A\n    SparseMatrix(indices=tensor([[0, 1, 2],\n                                 [3, 0, 2]]),\n                 values=tensor([40, 20, 60]),\n                 shape=(3, 4), nnz=3)\n    \"\"\"\n    # Python falls back to B.__radd__ then TypeError when NotImplemented is\n    # returned.\n    return spsp_add(A, B) if isinstance(B, SparseMatrix) else NotImplemented\n\n\ndef sp_sub(A: SparseMatrix, B: SparseMatrix) -> SparseMatrix:\n    \"\"\"Elementwise subtraction\n\n    Parameters\n    ----------\n    A : SparseMatrix\n        Sparse matrix\n    B : SparseMatrix\n        Sparse matrix\n\n    Returns\n    -------\n    SparseMatrix\n        Sparse matrix\n\n    Examples\n    --------\n\n    >>> indices = torch.tensor([[1, 0, 2], [0, 3, 2]])\n    >>> val = torch.tensor([10, 20, 30])\n    >>> val2 = torch.tensor([5, 10, 15])\n    >>> A = dglsp.spmatrix(indices, val, shape=(3, 4))\n    >>> B = dglsp.spmatrix(indices, val2, shape=(3, 4))\n    >>> A - B\n    SparseMatrix(indices=tensor([[0, 1, 2],\n                                 [3, 0, 2]]),\n                 values=tensor([10, 5, 15]),\n                 shape=(3, 4), nnz=3)\n    \"\"\"\n    # Python falls back to B.__rsub__ then TypeError when NotImplemented is\n    # returned.\n    return spsp_add(A, -B) if isinstance(B, SparseMatrix) else NotImplemented\n\n\ndef sp_mul(A: SparseMatrix, B: Union[SparseMatrix, Scalar]) -> SparseMatrix:\n    \"\"\"Elementwise multiplication\n\n    If :attr:`B` is a sparse matrix, both :attr:`A` and :attr:`B` must be\n    diagonal matrices.\n\n    Parameters\n    ----------\n    A : SparseMatrix\n        First operand\n    B : SparseMatrix or Scalar\n        Second operand\n\n    Returns\n    -------\n    SparseMatrix\n        Result of A * B\n\n    Examples\n    --------\n\n    >>> indices = torch.tensor([[1, 0, 2], [0, 3, 2]])\n    >>> val = torch.tensor([1, 2, 3])\n    >>> A = dglsp.spmatrix(indices, val, shape=(3, 4))\n\n    >>> A * 2\n    SparseMatrix(indices=tensor([[1, 0, 2],\n                                 [0, 3, 2]]),\n    values=tensor([2, 4, 6]),\n    shape=(3, 4), nnz=3)\n\n    >>> 2 * A\n    SparseMatrix(indices=tensor([[1, 0, 2],\n                                 [0, 3, 2]]),\n    values=tensor([2, 4, 6]),\n    shape=(3, 4), nnz=3)\n    \"\"\"\n    if is_scalar(B):\n        return val_like(A, A.val * B)\n    if A.is_diag() and B.is_diag():\n        assert A.shape == B.shape, (\n            f\"The shape of diagonal matrix A {A.shape} and B {B.shape} must\"\n            f\"match for elementwise multiplication.\"\n        )\n        return diag(A.val * B.val, A.shape)\n    # Python falls back to B.__rmul__(A) then TypeError when NotImplemented is\n    # returned.\n    # So this also handles the case of scalar * SparseMatrix since we set\n    # SparseMatrix.__rmul__ to be the same as SparseMatrix.__mul__.\n    return NotImplemented\n\n\ndef sp_div(A: SparseMatrix, B: Union[SparseMatrix, Scalar]) -> SparseMatrix:\n    \"\"\"Elementwise division\n\n    If :attr:`B` is a sparse matrix, both :attr:`A` and :attr:`B` must be\n    diagonal matrices.\n\n    Parameters\n    ----------\n    A : SparseMatrix\n        First operand\n    B : SparseMatrix or Scalar\n        Second operand\n\n    Returns\n    -------\n    SparseMatrix\n        Result of A / B\n\n    Examples\n    --------\n    >>> indices = torch.tensor([[1, 0, 2], [0, 3, 2]])\n    >>> val = torch.tensor([1, 2, 3])\n    >>> A = dglsp.spmatrix(indices, val, shape=(3, 4))\n    >>> A / 2\n    SparseMatrix(indices=tensor([[1, 0, 2],\n                                 [0, 3, 2]]),\n                 values=tensor([0.5000, 1.0000, 1.5000]),\n                 shape=(3, 4), nnz=3)\n    \"\"\"\n    if is_scalar(B):\n        return val_like(A, A.val / B)\n    if A.is_diag() and B.is_diag():\n        assert A.shape == B.shape, (\n            f\"The shape of diagonal matrix A {A.shape} and B {B.shape} must\"\n            f\"match for elementwise division.\"\n        )\n        return diag(A.val / B.val, A.shape)\n    # Python falls back to B.__rtruediv__(A) then TypeError when NotImplemented\n    # is returned.\n    return NotImplemented\n\n\ndef sp_power(A: SparseMatrix, scalar: Scalar) -> SparseMatrix:\n    \"\"\"Take the power of each nonzero element and return a sparse matrix with\n    the result.\n\n    Parameters\n    ----------\n    A : SparseMatrix\n        Sparse matrix\n    scalar : float or int\n        Exponent\n\n    Returns\n    -------\n    SparseMatrix\n        Sparse matrix\n\n    Examples\n    --------\n    >>> indices = torch.tensor([[1, 0, 2], [0, 3, 2]])\n    >>> val = torch.tensor([10, 20, 30])\n    >>> A = dglsp.spmatrix(indices, val)\n    >>> A ** 2\n    SparseMatrix(indices=tensor([[1, 0, 2],\n                                 [0, 3, 2]]),\n    values=tensor([100, 400, 900]),\n    shape=(3, 4), nnz=3)\n    \"\"\"\n    # Python falls back to scalar.__rpow__ then TypeError when NotImplemented\n    # is returned.\n    return val_like(A, A.val**scalar) if is_scalar(scalar) else NotImplemented\n\n\nSparseMatrix.__add__ = sp_add\nSparseMatrix.__sub__ = sp_sub\nSparseMatrix.__mul__ = sp_mul\nSparseMatrix.__rmul__ = sp_mul\nSparseMatrix.__truediv__ = sp_div\nSparseMatrix.__pow__ = sp_power\n", "path": "python/dgl/sparse/elementwise_op_sp.py"}], "after_files": [{"content": "\"\"\"DGL elementwise operators for sparse matrix module.\"\"\"\nfrom typing import Union\n\nimport torch\n\nfrom .sparse_matrix import diag, SparseMatrix, val_like\nfrom .utils import is_scalar, Scalar\n\n\ndef spsp_add(A, B):\n    \"\"\"Invoke C++ sparse library for addition\"\"\"\n    return SparseMatrix(\n        torch.ops.dgl_sparse.spsp_add(A.c_sparse_matrix, B.c_sparse_matrix)\n    )\n\n\ndef spsp_mul(A, B):\n    \"\"\"Invoke C++ sparse library for multiplication\"\"\"\n    return SparseMatrix(\n        torch.ops.dgl_sparse.spsp_mul(A.c_sparse_matrix, B.c_sparse_matrix)\n    )\n\n\ndef sp_add(A: SparseMatrix, B: SparseMatrix) -> SparseMatrix:\n    \"\"\"Elementwise addition\n\n    Parameters\n    ----------\n    A : SparseMatrix\n        Sparse matrix\n    B : SparseMatrix\n        Sparse matrix\n\n    Returns\n    -------\n    SparseMatrix\n        Sparse matrix\n\n    Examples\n    --------\n\n    >>> indices = torch.tensor([[1, 0, 2], [0, 3, 2]])\n    >>> val = torch.tensor([10, 20, 30])\n    >>> A = dglsp.spmatrix(indices, val, shape=(3, 4))\n    >>> A + A\n    SparseMatrix(indices=tensor([[0, 1, 2],\n                                 [3, 0, 2]]),\n                 values=tensor([40, 20, 60]),\n                 shape=(3, 4), nnz=3)\n    \"\"\"\n    # Python falls back to B.__radd__ then TypeError when NotImplemented is\n    # returned.\n    return spsp_add(A, B) if isinstance(B, SparseMatrix) else NotImplemented\n\n\ndef sp_sub(A: SparseMatrix, B: SparseMatrix) -> SparseMatrix:\n    \"\"\"Elementwise subtraction\n\n    Parameters\n    ----------\n    A : SparseMatrix\n        Sparse matrix\n    B : SparseMatrix\n        Sparse matrix\n\n    Returns\n    -------\n    SparseMatrix\n        Sparse matrix\n\n    Examples\n    --------\n\n    >>> indices = torch.tensor([[1, 0, 2], [0, 3, 2]])\n    >>> val = torch.tensor([10, 20, 30])\n    >>> val2 = torch.tensor([5, 10, 15])\n    >>> A = dglsp.spmatrix(indices, val, shape=(3, 4))\n    >>> B = dglsp.spmatrix(indices, val2, shape=(3, 4))\n    >>> A - B\n    SparseMatrix(indices=tensor([[0, 1, 2],\n                                 [3, 0, 2]]),\n                 values=tensor([10, 5, 15]),\n                 shape=(3, 4), nnz=3)\n    \"\"\"\n    # Python falls back to B.__rsub__ then TypeError when NotImplemented is\n    # returned.\n    return spsp_add(A, -B) if isinstance(B, SparseMatrix) else NotImplemented\n\n\ndef sp_mul(A: SparseMatrix, B: Union[SparseMatrix, Scalar]) -> SparseMatrix:\n    \"\"\"Elementwise multiplication\n\n    Note that if both :attr:`A` and :attr:`B` are sparse matrices, both of them\n    need to be diagonal or on CPU.\n\n    Parameters\n    ----------\n    A : SparseMatrix\n        First operand\n    B : SparseMatrix or Scalar\n        Second operand\n\n    Returns\n    -------\n    SparseMatrix\n        Result of A * B\n\n    Examples\n    --------\n\n    >>> indices = torch.tensor([[1, 0, 2], [0, 3, 2]])\n    >>> val = torch.tensor([1, 2, 3])\n    >>> A = dglsp.spmatrix(indices, val, shape=(3, 4))\n\n    >>> A * 2\n    SparseMatrix(indices=tensor([[1, 0, 2],\n                                 [0, 3, 2]]),\n                 values=tensor([2, 4, 6]),\n                 shape=(3, 4), nnz=3)\n\n    >>> 2 * A\n    SparseMatrix(indices=tensor([[1, 0, 2],\n                                 [0, 3, 2]]),\n                 values=tensor([2, 4, 6]),\n                 shape=(3, 4), nnz=3)\n\n    >>> indices2 = torch.tensor([[2, 0, 1], [0, 3, 2]])\n    >>> val2 = torch.tensor([3, 2, 1])\n    >>> B = dglsp.spmatrix(indices2, val2, shape=(3, 4))\n    >>> A * B\n    SparseMatrix(indices=tensor([[0],\n                                 [3]]),\n                 values=tensor([4]),\n                 shape=(3, 4), nnz=1)\n    \"\"\"\n    if is_scalar(B):\n        return val_like(A, A.val * B)\n    return spsp_mul(A, B)\n\n\ndef sp_div(A: SparseMatrix, B: Union[SparseMatrix, Scalar]) -> SparseMatrix:\n    \"\"\"Elementwise division\n\n    If :attr:`B` is a sparse matrix, both :attr:`A` and :attr:`B` must be\n    diagonal matrices.\n\n    Parameters\n    ----------\n    A : SparseMatrix\n        First operand\n    B : SparseMatrix or Scalar\n        Second operand\n\n    Returns\n    -------\n    SparseMatrix\n        Result of A / B\n\n    Examples\n    --------\n    >>> indices = torch.tensor([[1, 0, 2], [0, 3, 2]])\n    >>> val = torch.tensor([1, 2, 3])\n    >>> A = dglsp.spmatrix(indices, val, shape=(3, 4))\n    >>> A / 2\n    SparseMatrix(indices=tensor([[1, 0, 2],\n                                 [0, 3, 2]]),\n                 values=tensor([0.5000, 1.0000, 1.5000]),\n                 shape=(3, 4), nnz=3)\n    \"\"\"\n    if is_scalar(B):\n        return val_like(A, A.val / B)\n    if A.is_diag() and B.is_diag():\n        assert A.shape == B.shape, (\n            f\"The shape of diagonal matrix A {A.shape} and B {B.shape} must\"\n            f\"match for elementwise division.\"\n        )\n        return diag(A.val / B.val, A.shape)\n    # Python falls back to B.__rtruediv__(A) then TypeError when NotImplemented\n    # is returned.\n    return NotImplemented\n\n\ndef sp_power(A: SparseMatrix, scalar: Scalar) -> SparseMatrix:\n    \"\"\"Take the power of each nonzero element and return a sparse matrix with\n    the result.\n\n    Parameters\n    ----------\n    A : SparseMatrix\n        Sparse matrix\n    scalar : float or int\n        Exponent\n\n    Returns\n    -------\n    SparseMatrix\n        Sparse matrix\n\n    Examples\n    --------\n    >>> indices = torch.tensor([[1, 0, 2], [0, 3, 2]])\n    >>> val = torch.tensor([10, 20, 30])\n    >>> A = dglsp.spmatrix(indices, val)\n    >>> A ** 2\n    SparseMatrix(indices=tensor([[1, 0, 2],\n                                 [0, 3, 2]]),\n    values=tensor([100, 400, 900]),\n    shape=(3, 4), nnz=3)\n    \"\"\"\n    # Python falls back to scalar.__rpow__ then TypeError when NotImplemented\n    # is returned.\n    return val_like(A, A.val**scalar) if is_scalar(scalar) else NotImplemented\n\n\nSparseMatrix.__add__ = sp_add\nSparseMatrix.__sub__ = sp_sub\nSparseMatrix.__mul__ = sp_mul\nSparseMatrix.__rmul__ = sp_mul\nSparseMatrix.__truediv__ = sp_div\nSparseMatrix.__pow__ = sp_power\n", "path": "python/dgl/sparse/elementwise_op_sp.py"}]}