Repository: dmlc/dgl

Issue:
[Sparse] Support SparseMatrix element-wise division with the same sparsity.
## 🔨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/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
@@ -3,7 +3,7 @@
 
 import torch
 
-from .sparse_matrix import diag, SparseMatrix, val_like
+from .sparse_matrix import SparseMatrix, val_like
 from .utils import is_scalar, Scalar
 
 
@@ -21,6 +21,13 @@ def spsp_mul(A, B):
     )
 
 
+def spsp_div(A, B):
+    """Invoke C++ sparse library for division"""
+    return SparseMatrix(
+        torch.ops.dgl_sparse.spsp_div(A.c_sparse_matrix, B.c_sparse_matrix)
+    )
+
+
 def sp_add(A: SparseMatrix, B: SparseMatrix) -> SparseMatrix:
     """Elementwise addition
 
@@ -141,8 +148,9 @@ def sp_mul(A: SparseMatrix, B: Union[SparseMatrix, Scalar]) -> SparseMatrix:
 def sp_div(A: SparseMatrix, B: Union[SparseMatrix, Scalar]) -> SparseMatrix:
     """Elementwise division
 
-    If :attr:`B` is a sparse matrix, both :attr:`A` and :attr:`B` must be
-    diagonal matrices.
+    If :attr:`B` is a sparse matrix, both :attr:`A` and :attr:`B` must have the
+    same sparsity. And the returned matrix has the same order of non-zero
+    entries as :attr:`A`.
 
     Parameters
     ----------
@@ -169,15 +177,7 @@ def sp_div(A: SparseMatrix, B: Union[SparseMatrix, Scalar]) -> SparseMatrix:
     """
     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 division."
-        )
-        return diag(A.val / B.val, A.shape)
-    # Python falls back to B.__rtruediv__(A) then TypeError when NotImplemented
-    # is returned.
-    return NotImplemented
+    return spsp_div(A, B)
 
 
 def sp_power(A: SparseMatrix, scalar: Scalar) -> SparseMatrix:


Test patch:
diff --git a/tests/python/pytorch/sparse/test_elementwise_op.py b/tests/python/pytorch/sparse/test_elementwise_op.py
--- a/tests/python/pytorch/sparse/test_elementwise_op.py
+++ b/tests/python/pytorch/sparse/test_elementwise_op.py
@@ -225,7 +225,7 @@ def test_sub_sparse_diag(val_shape):
     assert torch.allclose(dense_diff, -diff4)
 
 
-@pytest.mark.parametrize("op", ["truediv", "pow"])
+@pytest.mark.parametrize("op", ["pow"])
 def test_error_op_sparse_diag(op):
     ctx = F.ctx()
     row = torch.tensor([1, 0, 2]).to(ctx)
diff --git a/tests/python/pytorch/sparse/test_elementwise_op_sp.py b/tests/python/pytorch/sparse/test_elementwise_op_sp.py
--- a/tests/python/pytorch/sparse/test_elementwise_op_sp.py
+++ b/tests/python/pytorch/sparse/test_elementwise_op_sp.py
@@ -4,7 +4,7 @@
 import pytest
 import torch
 
-from dgl.sparse import from_coo, mul, power, val_like
+from dgl.sparse import div, from_coo, mul, power, spmatrix, val_like
 
 from .utils import (
     rand_coo,
@@ -134,3 +134,24 @@ def test_spspmul(create_func1, create_func2, shape, nnz1, nnz2, nz_dim):
     assert torch.allclose(
         val_like(B, B.val.grad).to_dense(), DB.grad, atol=1e-05
     )
+
+
+@pytest.mark.parametrize(
+    "create_func", [rand_coo, rand_csr, rand_csc, rand_diag]
+)
+@pytest.mark.parametrize("shape", [(5, 5), (5, 3)])
+@pytest.mark.parametrize("nnz", [1, 14])
+@pytest.mark.parametrize("nz_dim", [None, 3])
+def test_spspdiv(create_func, nnz, shape, nz_dim):
+    dev = F.ctx()
+    A = create_func(shape, nnz, dev, nz_dim)
+
+    perm = torch.randperm(A.nnz, device=dev)
+    rperm = torch.argsort(perm)
+    B = spmatrix(A.indices()[:, perm], A.val[perm], A.shape)
+    C = div(A, B)
+    assert not C.has_duplicate()
+    assert torch.allclose(C.val, A.val / B.val[rperm], atol=1e-05)
+    assert torch.allclose(C.indices(), A.indices(), atol=1e-05)
+
+    # No need to test backward here, since it is handled by Pytorch
diff --git a/tests/python/pytorch/sparse/test_softmax.py b/tests/python/pytorch/sparse/test_softmax.py
--- a/tests/python/pytorch/sparse/test_softmax.py
+++ b/tests/python/pytorch/sparse/test_softmax.py
@@ -35,9 +35,9 @@ def test_softmax(val_D, csr, dim):
         g = dgl.graph((row, col), 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)
+    assert torch.allclose(A_max.val, score, atol=1e-05)
 
     grad = torch.randn_like(score).to(dev)
     A_max.val.backward(grad)
     score.backward(grad)
-    assert torch.allclose(A.val.grad, val_g.grad)
+    assert torch.allclose(A.val.grad, val_g.grad, atol=1e-05)