chore: import upstream snapshot with attribution

tests/python/pytorch/sparse/__init__.py

@@ -0,0 +1 @@
+""" DGL sparse tests"""

tests/python/pytorch/sparse/test_broadcast.py

@@ -0,0 +1,45 @@
+import operator
+
+import backend as F
+import pytest
+import torch
+
+from dgl.sparse import sp_broadcast_v
+
+from .utils import rand_coo
+
+
+@pytest.mark.parametrize("shape", [(3, 4), (1, 5), (5, 1)])
+@pytest.mark.parametrize("nnz", [1, 4])
+@pytest.mark.parametrize("nz_dim", [None, 2])
+@pytest.mark.parametrize("op", ["add", "sub", "mul", "truediv"])
+def test_sp_broadcast_v(shape, nnz, nz_dim, op):
+    dev = F.ctx()
+    A = rand_coo(shape, nnz, dev, nz_dim)
+
+    v = torch.randn(A.shape[1], device=dev)
+    res1 = sp_broadcast_v(A, v, op)
+    if A.val.dim() == 1:
+        rhs = v[A.col]
+    else:
+        rhs = v[A.col].view(-1, 1)
+    res2 = getattr(operator, op)(A.val, rhs)
+    assert torch.allclose(res1.val, res2)
+
+    v = torch.randn(1, A.shape[1], device=dev)
+    res1 = sp_broadcast_v(A, v, op)
+    if A.val.dim() == 1:
+        rhs = v.view(-1)[A.col]
+    else:
+        rhs = v.view(-1)[A.col].view(-1, 1)
+    res2 = getattr(operator, op)(A.val, rhs)
+    assert torch.allclose(res1.val, res2)
+
+    v = torch.randn(A.shape[0], 1, device=dev)
+    res1 = sp_broadcast_v(A, v, op)
+    if A.val.dim() == 1:
+        rhs = v.view(-1)[A.row]
+    else:
+        rhs = v.view(-1)[A.row].view(-1, 1)
+    res2 = getattr(operator, op)(A.val, rhs)
+    assert torch.allclose(res1.val, res2)

tests/python/pytorch/sparse/test_elementwise_op.py

@@ -0,0 +1,242 @@
+import operator
+
+import backend as F
+
+import dgl.sparse as dglsp
+import pytest
+import torch
+
+from dgl.sparse import diag, power
+
+
+@pytest.mark.parametrize("opname", ["add", "sub", "mul", "truediv"])
+def test_diag_op_diag(opname):
+    op = getattr(operator, opname)
+    ctx = F.ctx()
+    shape = (3, 4)
+    D1 = diag(torch.arange(1, 4).to(ctx), shape=shape)
+    D2 = diag(torch.arange(10, 13).to(ctx), shape=shape)
+    result = op(D1, D2)
+    assert torch.allclose(result.val, op(D1.val, D2.val), rtol=1e-4, atol=1e-4)
+    assert result.shape == D1.shape
+
+
+@pytest.mark.parametrize(
+    "v_scalar", [2, 2.5, torch.tensor(2), torch.tensor(2.5)]
+)
+def test_diag_op_scalar(v_scalar):
+    ctx = F.ctx()
+    shape = (3, 4)
+    D1 = diag(torch.arange(1, 4).to(ctx), shape=shape)
+
+    # D * v
+    D2 = D1 * v_scalar
+    assert torch.allclose(D1.val * v_scalar, D2.val, rtol=1e-4, atol=1e-4)
+    assert D1.shape == D2.shape
+
+    # v * D
+    D2 = v_scalar * D1
+    assert torch.allclose(v_scalar * D1.val, D2.val, rtol=1e-4, atol=1e-4)
+    assert D1.shape == D2.shape
+
+    # D / v
+    D2 = D1 / v_scalar
+    assert torch.allclose(D1.val / v_scalar, D2.val, rtol=1e-4, atol=1e-4)
+    assert D1.shape == D2.shape
+
+    # D ^ v
+    D1 = diag(torch.arange(1, 4).to(ctx))
+    D2 = D1**v_scalar
+    assert torch.allclose(D1.val**v_scalar, D2.val, rtol=1e-4, atol=1e-4)
+    assert D1.shape == D2.shape
+
+    # pow(D, v)
+    D2 = power(D1, v_scalar)
+    assert torch.allclose(D1.val**v_scalar, D2.val, rtol=1e-4, atol=1e-4)
+    assert D1.shape == D2.shape
+
+    with pytest.raises(TypeError):
+        D1 + v_scalar
+    with pytest.raises(TypeError):
+        v_scalar + D1
+
+    with pytest.raises(TypeError):
+        D1 - v_scalar
+    with pytest.raises(TypeError):
+        v_scalar - D1
+
+
+@pytest.mark.parametrize("val_shape", [(), (2,)])
+@pytest.mark.parametrize("opname", ["add", "sub"])
+def test_addsub_coo(val_shape, opname):
+    op = getattr(operator, opname)
+    func = getattr(dglsp, opname)
+    ctx = F.ctx()
+    row = torch.tensor([1, 0, 2]).to(ctx)
+    col = torch.tensor([0, 3, 2]).to(ctx)
+    val = torch.randn(row.shape + val_shape).to(ctx)
+    A = dglsp.from_coo(row, col, val)
+
+    row = torch.tensor([1, 0]).to(ctx)
+    col = torch.tensor([0, 2]).to(ctx)
+    val = torch.randn(row.shape + val_shape).to(ctx)
+    B = dglsp.from_coo(row, col, val, shape=A.shape)
+
+    C1 = op(A, B).to_dense()
+    C2 = func(A, B).to_dense()
+    dense_C = op(A.to_dense(), B.to_dense())
+
+    assert torch.allclose(dense_C, C1)
+    assert torch.allclose(dense_C, C2)
+
+    with pytest.raises(TypeError):
+        op(A, 2)
+    with pytest.raises(TypeError):
+        op(2, A)
+
+
+@pytest.mark.parametrize("val_shape", [(), (2,)])
+@pytest.mark.parametrize("opname", ["add", "sub"])
+def test_addsub_csr(val_shape, opname):
+    op = getattr(operator, opname)
+    func = getattr(dglsp, opname)
+    ctx = F.ctx()
+    indptr = torch.tensor([0, 1, 2, 3]).to(ctx)
+    indices = torch.tensor([3, 0, 2]).to(ctx)
+    val = torch.randn(indices.shape + val_shape).to(ctx)
+    A = dglsp.from_csr(indptr, indices, val)
+
+    indptr = torch.tensor([0, 1, 2, 2]).to(ctx)
+    indices = torch.tensor([2, 0]).to(ctx)
+    val = torch.randn(indices.shape + val_shape).to(ctx)
+    B = dglsp.from_csr(indptr, indices, val, shape=A.shape)
+
+    C1 = op(A, B).to_dense()
+    C2 = func(A, B).to_dense()
+    dense_C = op(A.to_dense(), B.to_dense())
+
+    assert torch.allclose(dense_C, C1)
+    assert torch.allclose(dense_C, C2)
+
+    with pytest.raises(TypeError):
+        op(A, 2)
+    with pytest.raises(TypeError):
+        op(2, A)
+
+
+@pytest.mark.parametrize("val_shape", [(), (2,)])
+@pytest.mark.parametrize("opname", ["add", "sub"])
+def test_addsub_csc(val_shape, opname):
+    op = getattr(operator, opname)
+    func = getattr(dglsp, opname)
+    ctx = F.ctx()
+    indptr = torch.tensor([0, 1, 1, 2, 3]).to(ctx)
+    indices = torch.tensor([1, 2, 0]).to(ctx)
+    val = torch.randn(indices.shape + val_shape).to(ctx)
+    A = dglsp.from_csc(indptr, indices, val)
+
+    indptr = torch.tensor([0, 1, 1, 2, 2]).to(ctx)
+    indices = torch.tensor([1, 0]).to(ctx)
+    val = torch.randn(indices.shape + val_shape).to(ctx)
+    B = dglsp.from_csc(indptr, indices, val, shape=A.shape)
+
+    C1 = op(A, B).to_dense()
+    C2 = func(A, B).to_dense()
+    dense_C = op(A.to_dense(), B.to_dense())
+
+    assert torch.allclose(dense_C, C1)
+    assert torch.allclose(dense_C, C2)
+
+    with pytest.raises(TypeError):
+        op(A, 2)
+    with pytest.raises(TypeError):
+        op(2, A)
+
+
+@pytest.mark.parametrize("val_shape", [(), (2,)])
+@pytest.mark.parametrize("opname", ["add", "sub"])
+def test_addsub_diag(val_shape, opname):
+    op = getattr(operator, opname)
+    func = getattr(dglsp, opname)
+    ctx = F.ctx()
+    shape = (3, 4)
+    val_shape = (shape[0],) + val_shape
+    D1 = dglsp.diag(torch.randn(val_shape).to(ctx), shape=shape)
+    D2 = dglsp.diag(torch.randn(val_shape).to(ctx), shape=shape)
+
+    C1 = op(D1, D2).to_dense()
+    C2 = func(D1, D2).to_dense()
+    dense_C = op(D1.to_dense(), D2.to_dense())
+
+    assert torch.allclose(dense_C, C1)
+    assert torch.allclose(dense_C, C2)
+
+    with pytest.raises(TypeError):
+        op(D1, 2)
+    with pytest.raises(TypeError):
+        op(2, D1)
+
+
+@pytest.mark.parametrize("val_shape", [(), (2,)])
+def test_add_sparse_diag(val_shape):
+    ctx = F.ctx()
+    row = torch.tensor([1, 0, 2]).to(ctx)
+    col = torch.tensor([0, 3, 2]).to(ctx)
+    val = torch.randn(row.shape + val_shape).to(ctx)
+    A = dglsp.from_coo(row, col, val)
+
+    shape = (3, 4)
+    val_shape = (shape[0],) + val_shape
+    D = dglsp.diag(torch.randn(val_shape).to(ctx), shape=shape)
+
+    sum1 = (A + D).to_dense()
+    sum2 = (D + A).to_dense()
+    sum3 = dglsp.add(A, D).to_dense()
+    sum4 = dglsp.add(D, A).to_dense()
+    dense_sum = A.to_dense() + D.to_dense()
+
+    assert torch.allclose(dense_sum, sum1)
+    assert torch.allclose(dense_sum, sum2)
+    assert torch.allclose(dense_sum, sum3)
+    assert torch.allclose(dense_sum, sum4)
+
+
+@pytest.mark.parametrize("val_shape", [(), (2,)])
+def test_sub_sparse_diag(val_shape):
+    ctx = F.ctx()
+    row = torch.tensor([1, 0, 2]).to(ctx)
+    col = torch.tensor([0, 3, 2]).to(ctx)
+    val = torch.randn(row.shape + val_shape).to(ctx)
+    A = dglsp.from_coo(row, col, val)
+
+    shape = (3, 4)
+    val_shape = (shape[0],) + val_shape
+    D = dglsp.diag(torch.randn(val_shape).to(ctx), shape=shape)
+
+    diff1 = (A - D).to_dense()
+    diff2 = (D - A).to_dense()
+    diff3 = dglsp.sub(A, D).to_dense()
+    diff4 = dglsp.sub(D, A).to_dense()
+    dense_diff = A.to_dense() - D.to_dense()
+
+    assert torch.allclose(dense_diff, diff1)
+    assert torch.allclose(dense_diff, -diff2)
+    assert torch.allclose(dense_diff, diff3)
+    assert torch.allclose(dense_diff, -diff4)
+
+
+@pytest.mark.parametrize("op", ["pow"])
+def test_error_op_sparse_diag(op):
+    ctx = F.ctx()
+    row = torch.tensor([1, 0, 2]).to(ctx)
+    col = torch.tensor([0, 3, 2]).to(ctx)
+    val = torch.randn(row.shape).to(ctx)
+    A = dglsp.from_coo(row, col, val)
+
+    shape = (3, 4)
+    D = dglsp.diag(torch.randn(row.shape[0]).to(ctx), shape=shape)
+
+    with pytest.raises(TypeError):
+        getattr(operator, op)(A, D)
+    with pytest.raises(TypeError):
+        getattr(operator, op)(D, A)

tests/python/pytorch/sparse/test_elementwise_op_sp.py

@@ -0,0 +1,157 @@
+import sys
+
+import backend as F
+import pytest
+import torch
+
+from dgl.sparse import div, from_coo, mul, power, spmatrix, val_like
+
+from .utils import (
+    rand_coo,
+    rand_csc,
+    rand_csr,
+    rand_diag,
+    sparse_matrix_to_dense,
+)
+
+
+def all_close_sparse(A, row, col, val, shape):
+    rowA, colA = A.coo()
+    valA = A.val
+    assert torch.allclose(rowA, row)
+    assert torch.allclose(colA, col)
+    assert torch.allclose(valA, val)
+    assert A.shape == shape
+
+
+@pytest.mark.parametrize(
+    "v_scalar", [2, 2.5, torch.tensor(2), torch.tensor(2.5)]
+)
+def test_muldiv_scalar(v_scalar):
+    ctx = F.ctx()
+    row = torch.tensor([1, 0, 2]).to(ctx)
+    col = torch.tensor([0, 3, 2]).to(ctx)
+    val = torch.randn(len(row)).to(ctx)
+    A1 = from_coo(row, col, val, shape=(3, 4))
+
+    # A * v
+    A2 = A1 * v_scalar
+    assert torch.allclose(A1.val * v_scalar, A2.val, rtol=1e-4, atol=1e-4)
+    assert A1.shape == A2.shape
+
+    # v * A
+    A2 = v_scalar * A1
+    assert torch.allclose(A1.val * v_scalar, A2.val, rtol=1e-4, atol=1e-4)
+    assert A1.shape == A2.shape
+
+    # A / v
+    A2 = A1 / v_scalar
+    assert torch.allclose(A1.val / v_scalar, A2.val, rtol=1e-4, atol=1e-4)
+    assert A1.shape == A2.shape
+
+    # v / A
+    with pytest.raises(TypeError):
+        v_scalar / A1
+
+
+@pytest.mark.parametrize("val_shape", [(3,), (3, 2)])
+def test_pow(val_shape):
+    # A ** v
+    ctx = F.ctx()
+    row = torch.tensor([1, 0, 2]).to(ctx)
+    col = torch.tensor([0, 3, 2]).to(ctx)
+    val = torch.randn(val_shape).to(ctx)
+    A = from_coo(row, col, val, shape=(3, 4))
+    exponent = 2
+    A_new = A**exponent
+    assert torch.allclose(A_new.val, val**exponent)
+    assert A_new.shape == A.shape
+    new_row, new_col = A_new.coo()
+    assert torch.allclose(new_row, row)
+    assert torch.allclose(new_col, col)
+
+    # power(A, v)
+    A_new = power(A, exponent)
+    assert torch.allclose(A_new.val, val**exponent)
+    assert A_new.shape == A.shape
+    new_row, new_col = A_new.coo()
+    assert torch.allclose(new_row, row)
+    assert torch.allclose(new_col, col)
+
+
+@pytest.mark.parametrize("op", ["add", "sub"])
+@pytest.mark.parametrize(
+    "v_scalar", [2, 2.5, torch.tensor(2), torch.tensor(2.5)]
+)
+def test_error_op_scalar(op, v_scalar):
+    ctx = F.ctx()
+    row = torch.tensor([1, 0, 2]).to(ctx)
+    col = torch.tensor([0, 3, 2]).to(ctx)
+    val = torch.randn(len(row)).to(ctx)
+    A = from_coo(row, col, val, shape=(3, 4))
+
+    with pytest.raises(TypeError):
+        A + v_scalar
+    with pytest.raises(TypeError):
+        v_scalar + A
+
+    with pytest.raises(TypeError):
+        A - v_scalar
+    with pytest.raises(TypeError):
+        v_scalar - A
+
+
+@pytest.mark.parametrize(
+    "create_func1", [rand_coo, rand_csr, rand_csc, rand_diag]
+)
+@pytest.mark.parametrize(
+    "create_func2", [rand_coo, rand_csr, rand_csc, rand_diag]
+)
+@pytest.mark.parametrize("shape", [(5, 5), (5, 3)])
+@pytest.mark.parametrize("nnz1", [5, 15])
+@pytest.mark.parametrize("nnz2", [1, 14])
+@pytest.mark.parametrize("nz_dim", [None, 3])
+def test_spspmul(create_func1, create_func2, shape, nnz1, nnz2, nz_dim):
+    dev = F.ctx()
+    A = create_func1(shape, nnz1, dev, nz_dim)
+    B = create_func2(shape, nnz2, dev, nz_dim)
+    C = mul(A, B)
+    assert not C.has_duplicate()
+
+    DA = sparse_matrix_to_dense(A)
+    DB = sparse_matrix_to_dense(B)
+    DC = DA * DB
+
+    grad = torch.rand_like(C.val)
+    C.val.backward(grad)
+    DC_grad = sparse_matrix_to_dense(val_like(C, grad))
+    DC.backward(DC_grad)
+
+    assert torch.allclose(sparse_matrix_to_dense(C), DC, atol=1e-05)
+    assert torch.allclose(
+        val_like(A, A.val.grad).to_dense(), DA.grad, atol=1e-05
+    )
+    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

tests/python/pytorch/sparse/test_matmul.py

@@ -0,0 +1,218 @@
+import warnings
+
+import backend as F
+import pytest
+import torch
+
+from dgl.sparse import bspmm, diag, from_coo, val_like
+from dgl.sparse.matmul import matmul
+
+from .utils import (
+    clone_detach_and_grad,
+    dense_mask,
+    rand_coo,
+    rand_csc,
+    rand_csr,
+    rand_stride,
+    sparse_matrix_to_dense,
+    sparse_matrix_to_torch_sparse,
+)
+
+
+def _torch_sparse_mm(torch_A1, torch_A2):
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", category=UserWarning)
+        return torch.sparse.mm(torch_A1, torch_A2)
+
+
+@pytest.mark.parametrize("create_func", [rand_coo, rand_csr, rand_csc])
+@pytest.mark.parametrize("shape", [(2, 7), (5, 2)])
+@pytest.mark.parametrize("nnz", [1, 10])
+@pytest.mark.parametrize("out_dim", [None, 10])
+def test_spmm(create_func, shape, nnz, out_dim):
+    dev = F.ctx()
+    A = create_func(shape, nnz, dev)
+    if out_dim is not None:
+        X = torch.randn(shape[1], out_dim, requires_grad=True, device=dev)
+    else:
+        X = torch.randn(shape[1], requires_grad=True, device=dev)
+
+    X = rand_stride(X)
+    sparse_result = matmul(A, X)
+    grad = torch.randn_like(sparse_result)
+    sparse_result.backward(grad)
+
+    adj = sparse_matrix_to_dense(A)
+    XX = clone_detach_and_grad(X)
+    dense_result = torch.matmul(adj, XX)
+    if out_dim is None:
+        dense_result = dense_result.view(-1)
+    dense_result.backward(grad)
+    assert torch.allclose(sparse_result, dense_result, atol=1e-05)
+    assert torch.allclose(X.grad, XX.grad, atol=1e-05)
+    assert torch.allclose(
+        dense_mask(adj.grad, A),
+        sparse_matrix_to_dense(val_like(A, A.val.grad)),
+        atol=1e-05,
+    )
+
+
+@pytest.mark.parametrize("create_func", [rand_coo, rand_csr, rand_csc])
+@pytest.mark.parametrize("shape", [(2, 7), (5, 2)])
+@pytest.mark.parametrize("nnz", [1, 10])
+def test_bspmm(create_func, shape, nnz):
+    dev = F.ctx()
+    A = create_func(shape, nnz, dev, 2)
+    X = torch.randn(shape[1], 10, 2, requires_grad=True, device=dev)
+    X = rand_stride(X)
+
+    sparse_result = matmul(A, X)
+    grad = torch.randn_like(sparse_result)
+    sparse_result.backward(grad)
+
+    XX = clone_detach_and_grad(X)
+    torch_A = A.to_dense().clone().detach().requires_grad_()
+    torch_result = torch_A.permute(2, 0, 1) @ XX.permute(2, 0, 1)
+
+    torch_result.backward(grad.permute(2, 0, 1))
+    assert torch.allclose(
+        sparse_result.permute(2, 0, 1), torch_result, atol=1e-05
+    )
+    assert torch.allclose(X.grad, XX.grad, atol=1e-05)
+    assert torch.allclose(
+        dense_mask(torch_A.grad, A),
+        sparse_matrix_to_dense(val_like(A, A.val.grad)),
+        atol=1e-05,
+    )
+
+
+@pytest.mark.parametrize("create_func1", [rand_coo, rand_csr, rand_csc])
+@pytest.mark.parametrize("create_func2", [rand_coo, rand_csr, rand_csc])
+@pytest.mark.parametrize("shape_n_m", [(5, 5), (5, 6)])
+@pytest.mark.parametrize("shape_k", [3, 4])
+@pytest.mark.parametrize("nnz1", [1, 10])
+@pytest.mark.parametrize("nnz2", [1, 10])
+def test_spspmm(create_func1, create_func2, shape_n_m, shape_k, nnz1, nnz2):
+    dev = F.ctx()
+    shape1 = shape_n_m
+    shape2 = (shape_n_m[1], shape_k)
+    A1 = create_func1(shape1, nnz1, dev)
+    A2 = create_func2(shape2, nnz2, dev)
+    A3 = matmul(A1, A2)
+    grad = torch.randn_like(A3.val)
+    A3.val.backward(grad)
+
+    torch_A1 = sparse_matrix_to_torch_sparse(A1)
+    torch_A2 = sparse_matrix_to_torch_sparse(A2)
+    torch_A3 = _torch_sparse_mm(torch_A1, torch_A2)
+    torch_A3_grad = sparse_matrix_to_torch_sparse(A3, grad)
+    torch_A3.backward(torch_A3_grad)
+
+    with torch.no_grad():
+        assert torch.allclose(A3.to_dense(), torch_A3.to_dense(), atol=1e-05)
+        assert torch.allclose(
+            val_like(A1, A1.val.grad).to_dense(),
+            torch_A1.grad.to_dense(),
+            atol=1e-05,
+        )
+        assert torch.allclose(
+            val_like(A2, A2.val.grad).to_dense(),
+            torch_A2.grad.to_dense(),
+            atol=1e-05,
+        )
+
+
+def test_spspmm_duplicate():
+    dev = F.ctx()
+
+    row = torch.tensor([1, 0, 0, 0, 1]).to(dev)
+    col = torch.tensor([1, 1, 1, 2, 2]).to(dev)
+    val = torch.randn(len(row)).to(dev)
+    shape = (4, 4)
+    A1 = from_coo(row, col, val, shape)
+
+    row = torch.tensor([1, 0, 0, 1]).to(dev)
+    col = torch.tensor([1, 1, 2, 2]).to(dev)
+    val = torch.randn(len(row)).to(dev)
+    shape = (4, 4)
+    A2 = from_coo(row, col, val, shape)
+
+    try:
+        matmul(A1, A2)
+    except:
+        pass
+    else:
+        assert False, "Should raise error."
+
+    try:
+        matmul(A2, A1)
+    except:
+        pass
+    else:
+        assert False, "Should raise error."
+
+
+@pytest.mark.parametrize("create_func", [rand_coo, rand_csr, rand_csc])
+@pytest.mark.parametrize("sparse_shape", [(5, 5), (5, 6)])
+@pytest.mark.parametrize("nnz", [1, 10])
+def test_sparse_diag_mm(create_func, sparse_shape, nnz):
+    dev = F.ctx()
+    diag_shape = sparse_shape[1], sparse_shape[1]
+    A = create_func(sparse_shape, nnz, dev)
+    diag_val = torch.randn(sparse_shape[1], device=dev, requires_grad=True)
+    D = diag(diag_val, diag_shape)
+    B = matmul(A, D)
+    grad = torch.randn_like(B.val)
+    B.val.backward(grad)
+
+    torch_A = sparse_matrix_to_torch_sparse(A)
+    torch_D = sparse_matrix_to_torch_sparse(D)
+    torch_B = _torch_sparse_mm(torch_A, torch_D)
+    torch_B_grad = sparse_matrix_to_torch_sparse(B, grad)
+    torch_B.backward(torch_B_grad)
+
+    with torch.no_grad():
+        assert torch.allclose(B.to_dense(), torch_B.to_dense(), atol=1e-05)
+        assert torch.allclose(
+            val_like(A, A.val.grad).to_dense(),
+            torch_A.grad.to_dense(),
+            atol=1e-05,
+        )
+        assert torch.allclose(
+            diag(D.val.grad, D.shape).to_dense(),
+            torch_D.grad.to_dense(),
+            atol=1e-05,
+        )
+
+
+@pytest.mark.parametrize("create_func", [rand_coo, rand_csr, rand_csc])
+@pytest.mark.parametrize("sparse_shape", [(5, 5), (5, 6)])
+@pytest.mark.parametrize("nnz", [1, 10])
+def test_diag_sparse_mm(create_func, sparse_shape, nnz):
+    dev = F.ctx()
+    diag_shape = sparse_shape[0], sparse_shape[0]
+    A = create_func(sparse_shape, nnz, dev)
+    diag_val = torch.randn(sparse_shape[0], device=dev, requires_grad=True)
+    D = diag(diag_val, diag_shape)
+    B = matmul(D, A)
+    grad = torch.randn_like(B.val)
+    B.val.backward(grad)
+
+    torch_A = sparse_matrix_to_torch_sparse(A)
+    torch_D = sparse_matrix_to_torch_sparse(D)
+    torch_B = _torch_sparse_mm(torch_D, torch_A)
+    torch_B_grad = sparse_matrix_to_torch_sparse(B, grad)
+    torch_B.backward(torch_B_grad)
+
+    with torch.no_grad():
+        assert torch.allclose(B.to_dense(), torch_B.to_dense(), atol=1e-05)
+        assert torch.allclose(
+            val_like(A, A.val.grad).to_dense(),
+            torch_A.grad.to_dense(),
+            atol=1e-05,
+        )
+        assert torch.allclose(
+            diag(D.val.grad, D.shape).to_dense(),
+            torch_D.grad.to_dense(),
+            atol=1e-05,
+        )

tests/python/pytorch/sparse/test_matrix_op.py

@@ -0,0 +1,47 @@
+import backend as F
+import pytest
+import torch
+
+from .utils import (
+    rand_coo,
+    rand_csc,
+    rand_csr,
+    rand_diag,
+    sparse_matrix_to_dense,
+)
+
+
+@pytest.mark.parametrize(
+    "create_func", [rand_diag, rand_csr, rand_csc, rand_coo]
+)
+@pytest.mark.parametrize("dim", [0, 1])
+@pytest.mark.parametrize("index", [None, (1, 3), (4, 0, 2)])
+def test_compact(create_func, dim, index):
+    ctx = F.ctx()
+    shape = (5, 5)
+    ans_idx = []
+    if index is not None:
+        ans_idx = list(dict.fromkeys(index))
+        index = torch.tensor(index).to(ctx)
+
+    A = create_func(shape, 8, ctx)
+
+    A_compact, ret_id = A.compact(dim, index)
+    A_compact_dense = sparse_matrix_to_dense(A_compact)
+
+    A_dense = sparse_matrix_to_dense(A)
+
+    for i in range(shape[dim]):
+        if dim == 0:
+            row = list(A_dense[i, :].nonzero().reshape(-1))
+        else:
+            row = list(A_dense[:, i].nonzero().reshape(-1))
+        if (i not in list(ans_idx)) and len(row) > 0:
+            ans_idx.append(i)
+    if len(ans_idx):
+        ans_idx = torch.tensor(ans_idx).to(ctx)
+    A_dense_select = sparse_matrix_to_dense(A.index_select(dim, ans_idx))
+
+    assert A_compact_dense.shape == A_dense_select.shape
+    assert torch.allclose(A_compact_dense, A_dense_select)
+    assert torch.allclose(ans_idx, ret_id)

tests/python/pytorch/sparse/test_reduction.py

@@ -0,0 +1,160 @@
+import doctest
+import operator
+import sys
+
+import backend as F
+
+import dgl.sparse as dglsp
+import pytest
+import torch
+
+dgl_op_map = {
+    "sum": "sum",
+    "amin": "smin",
+    "amax": "smax",
+    "mean": "smean",
+    "prod": "sprod",
+}
+default_entry = {
+    "sum": 0,
+    "amin": float("inf"),
+    "amax": float("-inf"),
+    "mean": 0,
+    "prod": 1,
+}
+binary_op_map = {
+    "sum": operator.add,
+    "amin": torch.min,
+    "amax": torch.max,
+    "mean": operator.add,
+    "prod": operator.mul,
+}
+
+NUM_ROWS = 10
+NUM_COLS = 15
+
+
+def _coalesce_dense(row, col, val, nrows, ncols, op):
+    # Sparse matrix coalescing on a dense matrix.
+    #
+    # It is done by stacking every non-zero entry on an individual slice
+    # of an (nrows x ncols x nnz), that is, construct a tensor A with
+    # shape (nrows, ncols, len(val)) where
+    #
+    #     A[row[i], col[i], i] = val[i]
+    #
+    # and then reducing on the third "nnz" dimension.
+    #
+    # The mask matrix M has the same sparsity pattern as A with 1 being
+    # the non-zero entries.  This is used for division if the reduce
+    # operator is mean.
+    M = torch.zeros(NUM_ROWS, NUM_COLS, device=F.ctx())
+    A = torch.full(
+        (NUM_ROWS, NUM_COLS, 20) + val.shape[1:],
+        default_entry[op],
+        device=F.ctx(),
+        dtype=val.dtype,
+    )
+    A = torch.index_put(A, (row, col, torch.arange(20)), val)
+    for i in range(20):
+        M[row[i], col[i]] += 1
+    if op == "mean":
+        A = A.sum(2)
+    else:
+        A = getattr(A, op)(2)
+    M = M.view(NUM_ROWS, NUM_COLS, *([1] * (val.dim() - 1)))
+    return A, M
+
+
+# Add docstring tests of dglsp.reduction to unit tests
+@pytest.mark.parametrize(
+    "func", ["reduce", "sum", "smin", "smax", "sprod", "smean"]
+)
+def test_docstring(func):
+    globs = {"torch": torch, "dglsp": dglsp}
+    runner = doctest.DebugRunner()
+    finder = doctest.DocTestFinder()
+    obj = getattr(dglsp, func)
+    for test in finder.find(obj, func, globs=globs):
+        runner.run(test)
+
+
+@pytest.mark.parametrize("shape", [(20,), (20, 20)])
+@pytest.mark.parametrize("op", ["sum", "amin", "amax", "mean", "prod"])
+@pytest.mark.parametrize("use_reduce", [False, True])
+def test_reduce_all(shape, op, use_reduce):
+    row = torch.randint(0, NUM_ROWS, (20,), device=F.ctx())
+    col = torch.randint(0, NUM_COLS, (20,), device=F.ctx())
+    val = torch.randn(*shape, device=F.ctx())
+    val2 = val.clone()
+    val = val.requires_grad_()
+    val2 = val2.requires_grad_()
+    A = dglsp.from_coo(row, col, val, shape=(NUM_ROWS, NUM_COLS))
+
+    A2, M = _coalesce_dense(row, col, val2, NUM_ROWS, NUM_COLS, op)
+
+    if not use_reduce:
+        output = getattr(A, dgl_op_map[op])()
+    else:
+        output = A.reduce(rtype=dgl_op_map[op])
+
+    if op == "mean":
+        output2 = A2.sum((0, 1)) / M.sum()
+    elif op == "prod":
+        output2 = A2.prod(0).prod(0)  # prod() does not support tuple of dims
+    else:
+        output2 = getattr(A2, op)((0, 1))
+    assert (output - output2).abs().max() < 1e-4
+
+    head = torch.randn(*output.shape).to(val) if output.dim() > 0 else None
+    output.backward(head)
+    output2.backward(head)
+    assert (val.grad - val2.grad).abs().max() < 1e-4
+
+
+@pytest.mark.parametrize("shape", [(20,), (20, 20)])
+@pytest.mark.parametrize("dim", [0, 1])
+@pytest.mark.parametrize("empty_nnz", [False, True])
+@pytest.mark.parametrize("op", ["sum", "amin", "amax", "mean", "prod"])
+@pytest.mark.parametrize("use_reduce", [False, True])
+def test_reduce_along(shape, dim, empty_nnz, op, use_reduce):
+    row = torch.randint(0, NUM_ROWS, (20,), device=F.ctx())
+    col = torch.randint(0, NUM_COLS, (20,), device=F.ctx())
+    if dim == 0:
+        mask = torch.bincount(col, minlength=NUM_COLS) == 0
+    else:
+        mask = torch.bincount(row, minlength=NUM_ROWS) == 0
+    val = torch.randn(*shape, device=F.ctx())
+    val2 = val.clone()
+    val = val.requires_grad_()
+    val2 = val2.requires_grad_()
+
+    # empty_nnz controls whether at least one column or one row has no
+    # non-zero entry.
+    if empty_nnz:
+        row[row == 0] = 1
+        col[col == 0] = 1
+
+    A = dglsp.from_coo(row, col, val, shape=(NUM_ROWS, NUM_COLS))
+
+    A2, M = _coalesce_dense(row, col, val2, NUM_ROWS, NUM_COLS, op)
+
+    if not use_reduce:
+        output = getattr(A, dgl_op_map[op])(dim)
+    else:
+        output = A.reduce(dim=dim, rtype=dgl_op_map[op])
+
+    if op == "mean":
+        output2 = A2.sum(dim) / M.sum(dim)
+    else:
+        output2 = getattr(A2, op)(dim)
+    zero_entry_idx = (M.sum(dim) != 0).nonzero(as_tuple=True)[0]
+    output3 = torch.index_put(
+        torch.zeros_like(output2), (zero_entry_idx,), output2[zero_entry_idx]
+    )
+    assert (output - output3).abs().max() < 1e-4
+
+    head = torch.randn(*output.shape).to(val) if output.dim() > 0 else None
+    output.backward(head)
+    output3.backward(head)
+    assert (val.grad - val2.grad).abs().max() < 1e-4

tests/python/pytorch/sparse/test_sddmm.py

@@ -0,0 +1,92 @@
+import sys
+
+import backend as F
+import pytest
+import torch
+
+from dgl.sparse import bsddmm, sddmm
+
+from .utils import (
+    clone_detach_and_grad,
+    rand_coo,
+    rand_csc,
+    rand_csr,
+    rand_stride,
+)
+
+
+@pytest.mark.parametrize("create_func", [rand_coo, rand_csr, rand_csc])
+@pytest.mark.parametrize("shape", [(5, 5), (5, 4)])
+@pytest.mark.parametrize("nnz", [2, 10])
+@pytest.mark.parametrize("hidden", [1, 5])
+def test_sddmm(create_func, shape, nnz, hidden):
+    dev = F.ctx()
+    A = create_func(shape, nnz, dev)
+    if hidden > 1:
+        B = torch.rand(shape[0], hidden, requires_grad=True, device=dev)
+        C = torch.rand(hidden, shape[1], requires_grad=True, device=dev)
+    else:
+        B = torch.rand(shape[0], requires_grad=True, device=dev)
+        C = torch.rand(shape[1], requires_grad=True, device=dev)
+
+    B = rand_stride(B)
+    C = rand_stride(C)
+
+    A_val_clone = clone_detach_and_grad(A.val)
+    dense_B = clone_detach_and_grad(B)
+    dense_C = clone_detach_and_grad(C)
+
+    sparse_result = sddmm(A, B, C)
+
+    grad = torch.rand_like(sparse_result.val)
+    sparse_result.val.backward(grad)
+
+    if hidden == 1:
+        dense_result = dense_B.view(-1, 1) @ dense_C.view(1, -1)
+    else:
+        dense_result = dense_B @ dense_C
+
+    row, col = A.coo()
+    dense_val = dense_result[row, col] * A_val_clone
+    dense_val.backward(grad)
+
+    assert torch.allclose(dense_val, sparse_result.val, atol=1e-05)
+    assert torch.allclose(dense_C.grad, C.grad, atol=1e-05)
+    assert torch.allclose(dense_B.grad, B.grad, atol=1e-05)
+    assert torch.allclose(A_val_clone.grad, A.val.grad, atol=1e-05)
+
+
+@pytest.mark.parametrize("create_func", [rand_coo, rand_csr, rand_csc])
+@pytest.mark.parametrize("shape", [(5, 5), (5, 4)])
+@pytest.mark.parametrize("nnz", [2, 10])
+@pytest.mark.parametrize("nz_dim", [2, 10])
+def test_bsddmm(create_func, shape, nnz, nz_dim):
+    dev = F.ctx()
+    hidden = 2
+    A = create_func(shape, nnz, dev, nz_dim)
+    B = torch.rand(shape[0], hidden, nz_dim, requires_grad=True, device=dev)
+    C = torch.rand(hidden, shape[1], nz_dim, requires_grad=True, device=dev)
+
+    B = rand_stride(B)
+    C = rand_stride(C)
+
+    A_val_clone = clone_detach_and_grad(A.val)
+    dense_B = clone_detach_and_grad(B)
+    dense_C = clone_detach_and_grad(C)
+
+    sparse_result = bsddmm(A, B, C)
+
+    grad = torch.rand_like(sparse_result.val)
+    sparse_result.val.backward(grad)
+
+    dense_result = dense_B.permute(2, 0, 1) @ dense_C.permute(2, 0, 1)
+    dense_result = dense_result.permute(1, 2, 0)
+
+    row, col = A.coo()
+    dense_val = dense_result[row, col] * A_val_clone
+    dense_val.backward(grad)
+
+    assert torch.allclose(dense_val, sparse_result.val, atol=1e-05)
+    assert torch.allclose(dense_C.grad, C.grad, atol=1e-05)
+    assert torch.allclose(dense_B.grad, B.grad, atol=1e-05)
+    assert torch.allclose(A_val_clone.grad, A.val.grad, atol=1e-05)

tests/python/pytorch/sparse/test_softmax.py

@@ -0,0 +1,43 @@
+import sys
+
+import backend as F
+
+import dgl
+import pytest
+import torch
+from dgl.sparse import from_coo, softmax
+
+
+@pytest.mark.parametrize("val_D", [None, 2])
+@pytest.mark.parametrize("csr", [True, False])
+@pytest.mark.parametrize("dim", [0, 1])
+def test_softmax(val_D, csr, dim):
+    dev = F.ctx()
+    row = torch.tensor([0, 0, 1, 1]).to(dev)
+    col = torch.tensor([0, 2, 1, 2]).to(dev)
+    nnz = len(row)
+    if val_D is None:
+        val = torch.randn(nnz).to(dev)
+    else:
+        val = torch.randn(nnz, val_D).to(dev)
+
+    val_sparse = val.clone().requires_grad_()
+    A = from_coo(row, col, val_sparse)
+
+    if csr:
+        # Test CSR
+        A.csr()
+
+    A_max = softmax(A, dim)
+    if dim == 1:
+        g = dgl.graph((col, row), num_nodes=max(A.shape))
+    else:
+        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, 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, atol=1e-05)

tests/python/pytorch/sparse/test_sparse_matrix.py

@@ -0,0 +1,878 @@
+import unittest
+import warnings
+
+import backend as F
+import pytest
+import torch
+
+from dgl.sparse import (
+    diag,
+    from_coo,
+    from_csc,
+    from_csr,
+    from_torch_sparse,
+    identity,
+    to_torch_sparse_coo,
+    to_torch_sparse_csc,
+    to_torch_sparse_csr,
+    val_like,
+)
+
+from .utils import (
+    rand_coo,
+    rand_csc,
+    rand_csr,
+    rand_diag,
+    sparse_matrix_to_dense,
+)
+
+
+def _torch_sparse_csr_tensor(indptr, indices, val, torch_sparse_shape):
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", category=UserWarning)
+        return torch.sparse_csr_tensor(indptr, indices, val, torch_sparse_shape)
+
+
+@pytest.mark.parametrize("dense_dim", [None, 4])
+@pytest.mark.parametrize("row", [(0, 0, 1, 2), (0, 1, 2, 4)])
+@pytest.mark.parametrize("col", [(0, 1, 2, 2), (1, 3, 3, 4)])
+@pytest.mark.parametrize("shape", [None, (5, 5), (5, 6)])
+def test_from_coo(dense_dim, row, col, shape):
+    val_shape = (len(row),)
+    if dense_dim is not None:
+        val_shape += (dense_dim,)
+    ctx = F.ctx()
+    val = torch.randn(val_shape).to(ctx)
+    row = torch.tensor(row).to(ctx)
+    col = torch.tensor(col).to(ctx)
+    mat = from_coo(row, col, val, shape)
+
+    if shape is None:
+        shape = (torch.max(row).item() + 1, torch.max(col).item() + 1)
+
+    mat_row, mat_col = mat.coo()
+    mat_val = mat.val
+
+    assert mat.shape == shape
+    assert mat.nnz == row.numel()
+    assert mat.dtype == val.dtype
+    assert torch.allclose(mat_val, val)
+    assert torch.allclose(mat_row, row)
+    assert torch.allclose(mat_col, col)
+
+
+@pytest.mark.parametrize("dense_dim", [None, 4])
+@pytest.mark.parametrize("indptr", [(0, 0, 1, 4), (0, 1, 2, 4)])
+@pytest.mark.parametrize("indices", [(0, 1, 2, 3), (1, 2, 3, 4)])
+@pytest.mark.parametrize("shape", [None, (3, 5)])
+def test_from_csr(dense_dim, indptr, indices, shape):
+    val_shape = (len(indices),)
+    if dense_dim is not None:
+        val_shape += (dense_dim,)
+    ctx = F.ctx()
+    val = torch.randn(val_shape).to(ctx)
+    indptr = torch.tensor(indptr).to(ctx)
+    indices = torch.tensor(indices).to(ctx)
+    mat = from_csr(indptr, indices, val, shape)
+
+    if shape is None:
+        shape = (indptr.numel() - 1, torch.max(indices).item() + 1)
+
+    assert mat.device == val.device
+    assert mat.shape == shape
+    assert mat.nnz == indices.numel()
+    assert mat.dtype == val.dtype
+    mat_indptr, mat_indices, value_indices = mat.csr()
+    mat_val = mat.val if value_indices is None else mat.val[value_indices]
+    assert torch.allclose(mat_indptr, indptr)
+    assert torch.allclose(mat_indices, indices)
+    assert torch.allclose(mat_val, val)
+
+
+@pytest.mark.parametrize("dense_dim", [None, 4])
+@pytest.mark.parametrize("indptr", [(0, 0, 1, 4), (0, 1, 2, 4)])
+@pytest.mark.parametrize("indices", [(0, 1, 2, 3), (1, 2, 3, 4)])
+@pytest.mark.parametrize("shape", [None, (5, 3)])
+def test_from_csc(dense_dim, indptr, indices, shape):
+    val_shape = (len(indices),)
+    if dense_dim is not None:
+        val_shape += (dense_dim,)
+    ctx = F.ctx()
+    val = torch.randn(val_shape).to(ctx)
+    indptr = torch.tensor(indptr).to(ctx)
+    indices = torch.tensor(indices).to(ctx)
+    mat = from_csc(indptr, indices, val, shape)
+
+    if shape is None:
+        shape = (torch.max(indices).item() + 1, indptr.numel() - 1)
+
+    assert mat.device == val.device
+    assert mat.shape == shape
+    assert mat.nnz == indices.numel()
+    assert mat.dtype == val.dtype
+    mat_indptr, mat_indices, value_indices = mat.csc()
+    mat_val = mat.val if value_indices is None else mat.val[value_indices]
+    assert torch.allclose(mat_indptr, indptr)
+    assert torch.allclose(mat_indices, indices)
+    assert torch.allclose(mat_val, val)
+
+
+@pytest.mark.parametrize("val_shape", [(3), (3, 2)])
+def test_dense(val_shape):
+    ctx = F.ctx()
+
+    row = torch.tensor([1, 1, 2]).to(ctx)
+    col = torch.tensor([2, 4, 3]).to(ctx)
+    val = torch.randn(val_shape).to(ctx)
+    A = from_coo(row, col, val)
+    A_dense = A.to_dense()
+
+    shape = A.shape + val.shape[1:]
+    mat = torch.zeros(shape, device=ctx)
+    mat[row, col] = val
+    assert torch.allclose(A_dense, mat)
+
+
+@pytest.mark.parametrize("dense_dim", [None, 4])
+@pytest.mark.parametrize("indptr", [(0, 0, 1, 4), (0, 1, 2, 4)])
+@pytest.mark.parametrize("indices", [(0, 1, 2, 3), (1, 4, 3, 2)])
+@pytest.mark.parametrize("shape", [None, (3, 5)])
+def test_csr_to_coo(dense_dim, indptr, indices, shape):
+    ctx = F.ctx()
+    val_shape = (len(indices),)
+    if dense_dim is not None:
+        val_shape += (dense_dim,)
+    val = torch.randn(val_shape).to(ctx)
+    indptr = torch.tensor(indptr).to(ctx)
+    indices = torch.tensor(indices).to(ctx)
+    mat = from_csr(indptr, indices, val, shape)
+
+    if shape is None:
+        shape = (indptr.numel() - 1, torch.max(indices).item() + 1)
+
+    row = (
+        torch.arange(0, indptr.shape[0] - 1)
+        .to(ctx)
+        .repeat_interleave(torch.diff(indptr))
+    )
+    col = indices
+    mat_row, mat_col = mat.coo()
+    mat_val = mat.val
+
+    assert mat.shape == shape
+    assert mat.nnz == row.numel()
+    assert mat.device == row.device
+    assert mat.dtype == val.dtype
+    assert torch.allclose(mat_val, val)
+    assert torch.allclose(mat_row, row)
+    assert torch.allclose(mat_col, col)
+
+
+@pytest.mark.parametrize("dense_dim", [None, 4])
+@pytest.mark.parametrize("indptr", [(0, 0, 1, 4), (0, 1, 2, 4)])
+@pytest.mark.parametrize("indices", [(0, 1, 2, 3), (1, 4, 3, 2)])
+@pytest.mark.parametrize("shape", [None, (5, 3)])
+def test_csc_to_coo(dense_dim, indptr, indices, shape):
+    ctx = F.ctx()
+    val_shape = (len(indices),)
+    if dense_dim is not None:
+        val_shape += (dense_dim,)
+    val = torch.randn(val_shape).to(ctx)
+    indptr = torch.tensor(indptr).to(ctx)
+    indices = torch.tensor(indices).to(ctx)
+    mat = from_csc(indptr, indices, val, shape)
+
+    if shape is None:
+        shape = (torch.max(indices).item() + 1, indptr.numel() - 1)
+
+    col = (
+        torch.arange(0, indptr.shape[0] - 1)
+        .to(ctx)
+        .repeat_interleave(torch.diff(indptr))
+    )
+    row = indices
+    mat_row, mat_col = mat.coo()
+    mat_val = mat.val
+
+    assert mat.shape == shape
+    assert mat.nnz == row.numel()
+    assert mat.device == row.device
+    assert mat.dtype == val.dtype
+    assert torch.allclose(mat_val, val)
+    assert torch.allclose(mat_row, row)
+    assert torch.allclose(mat_col, col)
+
+
+def _scatter_add(a, index, v=1):
+    index = index.tolist()
+    for i in index:
+        a[i] += v
+    return a
+
+
+@pytest.mark.parametrize("dense_dim", [None, 4])
+@pytest.mark.parametrize("row", [(0, 0, 1, 2), (0, 1, 2, 4)])
+@pytest.mark.parametrize("col", [(0, 1, 2, 2), (1, 3, 3, 4)])
+@pytest.mark.parametrize("shape", [None, (5, 5), (5, 6)])
+def test_coo_to_csr(dense_dim, row, col, shape):
+    val_shape = (len(row),)
+    if dense_dim is not None:
+        val_shape += (dense_dim,)
+    ctx = F.ctx()
+    val = torch.randn(val_shape).to(ctx)
+    row = torch.tensor(row).to(ctx)
+    col = torch.tensor(col).to(ctx)
+    mat = from_coo(row, col, val, shape)
+
+    if shape is None:
+        shape = (torch.max(row).item() + 1, torch.max(col).item() + 1)
+
+    mat_indptr, mat_indices, value_indices = mat.csr()
+    mat_val = mat.val if value_indices is None else mat.val[value_indices]
+    indptr = torch.zeros(shape[0] + 1).to(ctx)
+    indptr = _scatter_add(indptr, row + 1)
+    indptr = torch.cumsum(indptr, 0).long()
+    indices = col
+
+    assert mat.shape == shape
+    assert mat.nnz == row.numel()
+    assert mat.dtype == val.dtype
+    assert torch.allclose(mat_val, val)
+    assert torch.allclose(mat_indptr, indptr)
+    assert torch.allclose(mat_indices, indices)
+
+
+@pytest.mark.parametrize("dense_dim", [None, 4])
+@pytest.mark.parametrize("indptr", [(0, 0, 1, 4), (0, 1, 2, 4)])
+@pytest.mark.parametrize("indices", [(0, 1, 2, 3), (1, 4, 3, 2)])
+@pytest.mark.parametrize("shape", [None, (5, 3)])
+def test_csc_to_csr(dense_dim, indptr, indices, shape):
+    ctx = F.ctx()
+    val_shape = (len(indices),)
+    if dense_dim is not None:
+        val_shape += (dense_dim,)
+    val = torch.randn(val_shape).to(ctx)
+    indptr = torch.tensor(indptr).to(ctx)
+    indices = torch.tensor(indices).to(ctx)
+    mat = from_csc(indptr, indices, val, shape)
+    mat_indptr, mat_indices, value_indices = mat.csr()
+    mat_val = mat.val if value_indices is None else mat.val[value_indices]
+
+    if shape is None:
+        shape = (torch.max(indices).item() + 1, indptr.numel() - 1)
+
+    col = (
+        torch.arange(0, indptr.shape[0] - 1)
+        .to(ctx)
+        .repeat_interleave(torch.diff(indptr))
+    )
+    row = indices
+    row, sort_index = row.sort(stable=True)
+    col = col[sort_index]
+    val = val[sort_index]
+    indptr = torch.zeros(shape[0] + 1).to(ctx)
+    indptr = _scatter_add(indptr, row + 1)
+    indptr = torch.cumsum(indptr, 0).long()
+    indices = col
+
+    assert mat.shape == shape
+    assert mat.nnz == row.numel()
+    assert mat.device == row.device
+    assert mat.dtype == val.dtype
+    assert torch.allclose(mat_val, val)
+    assert torch.allclose(mat_indptr, indptr)
+    assert torch.allclose(mat_indices, indices)
+
+
+@pytest.mark.parametrize("dense_dim", [None, 4])
+@pytest.mark.parametrize("row", [(0, 0, 1, 2), (0, 1, 2, 4)])
+@pytest.mark.parametrize("col", [(0, 1, 2, 2), (1, 3, 3, 4)])
+@pytest.mark.parametrize("shape", [None, (5, 5), (5, 6)])
+def test_coo_to_csc(dense_dim, row, col, shape):
+    val_shape = (len(row),)
+    if dense_dim is not None:
+        val_shape += (dense_dim,)
+    ctx = F.ctx()
+    val = torch.randn(val_shape).to(ctx)
+    row = torch.tensor(row).to(ctx)
+    col = torch.tensor(col).to(ctx)
+    mat = from_coo(row, col, val, shape)
+
+    if shape is None:
+        shape = (torch.max(row).item() + 1, torch.max(col).item() + 1)
+
+    mat_indptr, mat_indices, value_indices = mat.csc()
+    mat_val = mat.val if value_indices is None else mat.val[value_indices]
+    indptr = torch.zeros(shape[1] + 1).to(ctx)
+    _scatter_add(indptr, col + 1)
+    indptr = torch.cumsum(indptr, 0).long()
+    indices = row
+
+    assert mat.shape == shape
+    assert mat.nnz == row.numel()
+    assert mat.dtype == val.dtype
+    assert torch.allclose(mat_val, val)
+    assert torch.allclose(mat_indptr, indptr)
+    assert torch.allclose(mat_indices, indices)
+
+
+@pytest.mark.parametrize("dense_dim", [None, 4])
+@pytest.mark.parametrize("indptr", [(0, 0, 1, 4), (0, 1, 2, 4)])
+@pytest.mark.parametrize("indices", [(0, 1, 2, 3), (1, 2, 3, 4)])
+@pytest.mark.parametrize("shape", [None, (3, 5)])
+def test_csr_to_csc(dense_dim, indptr, indices, shape):
+    val_shape = (len(indices),)
+    if dense_dim is not None:
+        val_shape += (dense_dim,)
+    ctx = F.ctx()
+    val = torch.randn(val_shape).to(ctx)
+    indptr = torch.tensor(indptr).to(ctx)
+    indices = torch.tensor(indices).to(ctx)
+    mat = from_csr(indptr, indices, val, shape)
+    mat_indptr, mat_indices, value_indices = mat.csc()
+    mat_val = mat.val if value_indices is None else mat.val[value_indices]
+
+    if shape is None:
+        shape = (indptr.numel() - 1, torch.max(indices).item() + 1)
+
+    row = (
+        torch.arange(0, indptr.shape[0] - 1)
+        .to(ctx)
+        .repeat_interleave(torch.diff(indptr))
+    )
+
+    col = indices
+    col, sort_index = col.sort(stable=True)
+    row = row[sort_index]
+    val = val[sort_index]
+    indptr = torch.zeros(shape[1] + 1).to(ctx)
+    indptr = _scatter_add(indptr, col + 1)
+    indptr = torch.cumsum(indptr, 0).long()
+    indices = row
+
+    assert mat.shape == shape
+    assert mat.nnz == row.numel()
+    assert mat.device == row.device
+    assert mat.dtype == val.dtype
+    assert torch.allclose(mat_val, val)
+    assert torch.allclose(mat_indptr, indptr)
+    assert torch.allclose(mat_indices, indices)
+
+
+@pytest.mark.parametrize("shape", [(3, 5), (5, 5), (5, 4)])
+def test_diag_conversions(shape):
+    n_rows, n_cols = shape
+    nnz = min(shape)
+    ctx = F.ctx()
+    val = torch.randn(nnz).to(ctx)
+    D = diag(val, shape)
+    row, col = D.coo()
+    assert torch.allclose(row, torch.arange(nnz).to(ctx))
+    assert torch.allclose(col, torch.arange(nnz).to(ctx))
+
+    indptr, indices, _ = D.csr()
+    exp_indptr = list(range(0, nnz + 1)) + [nnz] * (n_rows - nnz)
+    assert torch.allclose(indptr, torch.tensor(exp_indptr).to(ctx))
+    assert torch.allclose(indices, torch.arange(nnz).to(ctx))
+
+    indptr, indices, _ = D.csc()
+    exp_indptr = list(range(0, nnz + 1)) + [nnz] * (n_cols - nnz)
+    assert torch.allclose(indptr, torch.tensor(exp_indptr).to(ctx))
+    assert torch.allclose(indices, torch.arange(nnz).to(ctx))
+
+
+@pytest.mark.parametrize("val_shape", [(3), (3, 2)])
+@pytest.mark.parametrize("shape", [(3, 5), (5, 5)])
+def test_val_like(val_shape, shape):
+    def check_val_like(A, B):
+        assert A.shape == B.shape
+        assert A.nnz == B.nnz
+        assert torch.allclose(torch.stack(A.coo()), torch.stack(B.coo()))
+        assert A.val.device == B.val.device
+
+    ctx = F.ctx()
+
+    # COO
+    row = torch.tensor([1, 1, 2]).to(ctx)
+    col = torch.tensor([2, 4, 3]).to(ctx)
+    val = torch.randn(3).to(ctx)
+    coo_A = from_coo(row, col, val, shape)
+    new_val = torch.randn(val_shape).to(ctx)
+    coo_B = val_like(coo_A, new_val)
+    check_val_like(coo_A, coo_B)
+
+    # CSR
+    indptr, indices, _ = coo_A.csr()
+    csr_A = from_csr(indptr, indices, val, shape)
+    csr_B = val_like(csr_A, new_val)
+    check_val_like(csr_A, csr_B)
+
+    # CSC
+    indptr, indices, _ = coo_A.csc()
+    csc_A = from_csc(indptr, indices, val, shape)
+    csc_B = val_like(csc_A, new_val)
+    check_val_like(csc_A, csc_B)
+
+
+def test_coalesce():
+    ctx = F.ctx()
+
+    row = torch.tensor([1, 0, 0, 0, 1]).to(ctx)
+    col = torch.tensor([1, 1, 1, 2, 2]).to(ctx)
+    val = torch.arange(len(row)).to(ctx)
+    A = from_coo(row, col, val, (4, 4))
+
+    assert A.has_duplicate()
+
+    A_coalesced = A.coalesce()
+
+    assert A_coalesced.nnz == 4
+    assert A_coalesced.shape == (4, 4)
+    assert list(A_coalesced.row) == [0, 0, 1, 1]
+    assert list(A_coalesced.col) == [1, 2, 1, 2]
+    # Values of duplicate indices are added together.
+    assert list(A_coalesced.val) == [3, 3, 0, 4]
+    assert not A_coalesced.has_duplicate()
+
+
+def test_has_duplicate():
+    ctx = F.ctx()
+
+    row = torch.tensor([1, 0, 0, 0, 1]).to(ctx)
+    col = torch.tensor([1, 1, 1, 2, 2]).to(ctx)
+    val = torch.arange(len(row)).to(ctx)
+    shape = (4, 4)
+
+    # COO
+    coo_A = from_coo(row, col, val, shape)
+    assert coo_A.has_duplicate()
+
+    # CSR
+    indptr, indices, _ = coo_A.csr()
+    csr_A = from_csr(indptr, indices, val, shape)
+    assert csr_A.has_duplicate()
+
+    # CSC
+    indptr, indices, _ = coo_A.csc()
+    csc_A = from_csc(indptr, indices, val, shape)
+    assert csc_A.has_duplicate()
+
+
+@pytest.mark.parametrize(
+    "create_func", [rand_diag, rand_csr, rand_csc, rand_coo]
+)
+@pytest.mark.parametrize("shape", [(5, 5), (6, 4)])
+@pytest.mark.parametrize("dense_dim", [None, 4])
+@pytest.mark.parametrize("select_dim", [0, 1])
+@pytest.mark.parametrize("index", [(0, 1, 3), (1, 2)])
+def test_index_select(create_func, shape, dense_dim, select_dim, index):
+    ctx = F.ctx()
+    A = create_func(shape, 20, ctx, dense_dim)
+    index = torch.tensor(index).to(ctx)
+    A_select = A.index_select(select_dim, index)
+
+    dense = sparse_matrix_to_dense(A)
+    dense_select = torch.index_select(dense, select_dim, index)
+
+    A_select_to_dense = sparse_matrix_to_dense(A_select)
+
+    assert A_select_to_dense.shape == dense_select.shape
+    assert torch.allclose(A_select_to_dense, dense_select)
+
+
+@pytest.mark.parametrize(
+    "create_func", [rand_diag, rand_csr, rand_csc, rand_coo]
+)
+@pytest.mark.parametrize("shape", [(5, 5), (6, 4)])
+@pytest.mark.parametrize("dense_dim", [None, 4])
+@pytest.mark.parametrize("select_dim", [0, 1])
+@pytest.mark.parametrize("rang", [slice(0, 2), slice(1, 3)])
+def test_range_select(create_func, shape, dense_dim, select_dim, rang):
+    ctx = F.ctx()
+    A = create_func(shape, 20, ctx, dense_dim)
+    A_select = A.range_select(select_dim, rang)
+
+    dense = sparse_matrix_to_dense(A)
+    if select_dim == 0:
+        dense_select = dense[rang, :]
+    else:
+        dense_select = dense[:, rang]
+
+    A_select_to_dense = sparse_matrix_to_dense(A_select)
+
+    assert A_select_to_dense.shape == dense_select.shape
+    assert torch.allclose(A_select_to_dense, dense_select)
+
+
+@pytest.mark.parametrize(
+    "create_func", [rand_diag, rand_csr, rand_csc, rand_coo]
+)
+@pytest.mark.parametrize("index", [(0, 1, 2, 3, 4), (0, 1, 3), (1, 1, 2)])
+@pytest.mark.parametrize("replace", [False, True])
+@pytest.mark.parametrize("bias", [False, True])
+def test_sample_rowwise(create_func, index, replace, bias):
+    ctx = F.ctx()
+    shape = (5, 5)
+    sample_dim = 0
+    sample_num = 3
+    A = create_func(shape, 10, ctx)
+    A = val_like(A, torch.abs(A.val))
+
+    index = torch.tensor(index).to(ctx)
+
+    A_sample = A.sample(sample_dim, sample_num, index, replace, bias)
+    A_dense = sparse_matrix_to_dense(A)
+    A_sample_to_dense = sparse_matrix_to_dense(A_sample)
+
+    ans_shape = (index.size(0), shape[1])
+    # Verify sample elements in origin rows
+    for i, row in enumerate(list(index)):
+        ans_ele = list(A_dense[row, :].nonzero().reshape(-1))
+        ret_ele = list(A_sample_to_dense[i, :].nonzero().reshape(-1))
+        for e in ret_ele:
+            assert e in ans_ele
+        if replace:
+            # The number of sample elements in one row should be equal to
+            # 'sample_num' if the row is not empty otherwise should be
+            # equal to 0.
+            assert list(A_sample.row).count(torch.tensor(i)) == (
+                sample_num if len(ans_ele) != 0 else 0
+            )
+        else:
+            assert len(ret_ele) == min(sample_num, len(ans_ele))
+
+    assert A_sample.shape == ans_shape
+    if not replace:
+        assert not A_sample.has_duplicate()
+
+
+@pytest.mark.parametrize(
+    "create_func", [rand_diag, rand_csr, rand_csc, rand_coo]
+)
+@pytest.mark.parametrize("index", [(0, 1, 2, 3, 4), (0, 1, 3), (1, 1, 2)])
+@pytest.mark.parametrize("replace", [False, True])
+@pytest.mark.parametrize("bias", [False, True])
+def test_sample_columnwise(create_func, index, replace, bias):
+    ctx = F.ctx()
+    shape = (5, 5)
+    sample_dim = 1
+    sample_num = 3
+    A = create_func(shape, 10, ctx)
+    A = val_like(A, torch.abs(A.val))
+
+    index = torch.tensor(index).to(ctx)
+
+    A_sample = A.sample(sample_dim, sample_num, index, replace, bias)
+    A_dense = sparse_matrix_to_dense(A)
+    A_sample_to_dense = sparse_matrix_to_dense(A_sample)
+
+    ans_shape = (shape[0], index.size(0))
+    # Verify sample elements in origin columns
+    for i, col in enumerate(list(index)):
+        ans_ele = list(A_dense[:, col].nonzero().reshape(-1))
+        ret_ele = list(A_sample_to_dense[:, i].nonzero().reshape(-1))
+        for e in ret_ele:
+            assert e in ans_ele
+        if replace:
+            # The number of sample elements in one column should be equal to
+            # 'sample_num' if the column is not empty otherwise should be
+            # equal to 0.
+            assert list(A_sample.col).count(torch.tensor(i)) == (
+                sample_num if len(ans_ele) != 0 else 0
+            )
+        else:
+            assert len(ret_ele) == min(sample_num, len(ans_ele))
+
+    assert A_sample.shape == ans_shape
+    if not replace:
+        assert not A_sample.has_duplicate()
+
+
+def test_print():
+    ctx = F.ctx()
+
+    # basic
+    row = torch.tensor([1, 1, 3]).to(ctx)
+    col = torch.tensor([2, 1, 3]).to(ctx)
+    val = torch.tensor([1.0, 1.0, 2.0]).to(ctx)
+    A = from_coo(row, col, val)
+    expected = (
+        str(
+            """SparseMatrix(indices=tensor([[1, 1, 3],
+                             [2, 1, 3]]),
+             values=tensor([1., 1., 2.]),
+             shape=(4, 4), nnz=3)"""
+        )
+        if str(ctx) == "cpu"
+        else str(
+            """SparseMatrix(indices=tensor([[1, 1, 3],
+                             [2, 1, 3]], device='cuda:0'),
+             values=tensor([1., 1., 2.], device='cuda:0'),
+             shape=(4, 4), nnz=3)"""
+        )
+    )
+    assert str(A) == expected, print(A, expected)
+
+    # vector-shape non zero
+    row = torch.tensor([1, 1, 3]).to(ctx)
+    col = torch.tensor([2, 1, 3]).to(ctx)
+    val = torch.tensor(
+        [[1.3080, 1.5984], [-0.4126, 0.7250], [-0.5416, -0.7022]]
+    ).to(ctx)
+    A = from_coo(row, col, val)
+    expected = (
+        str(
+            """SparseMatrix(indices=tensor([[1, 1, 3],
+                             [2, 1, 3]]),
+             values=tensor([[ 1.3080,  1.5984],
+                            [-0.4126,  0.7250],
+                            [-0.5416, -0.7022]]),
+             shape=(4, 4), nnz=3, val_size=(2,))"""
+        )
+        if str(ctx) == "cpu"
+        else str(
+            """SparseMatrix(indices=tensor([[1, 1, 3],
+                             [2, 1, 3]], device='cuda:0'),
+             values=tensor([[ 1.3080,  1.5984],
+                            [-0.4126,  0.7250],
+                            [-0.5416, -0.7022]], device='cuda:0'),
+             shape=(4, 4), nnz=3, val_size=(2,))"""
+        )
+    )
+    assert str(A) == expected, print(A, expected)
+
+
+@unittest.skipIf(
+    F._default_context_str == "cpu",
+    reason="Device conversions don't need to be tested on CPU.",
+)
+@pytest.mark.parametrize("device", ["cpu", "cuda"])
+def test_to_device(device):
+    row = torch.tensor([1, 1, 2])
+    col = torch.tensor([1, 2, 0])
+    mat = from_coo(row, col, shape=(3, 4))
+
+    target_row = row.to(device)
+    target_col = col.to(device)
+    target_val = mat.val.to(device)
+
+    mat2 = mat.to(device=device)
+    assert mat2.shape == mat.shape
+    assert torch.allclose(mat2.row, target_row)
+    assert torch.allclose(mat2.col, target_col)
+    assert torch.allclose(mat2.val, target_val)
+
+    mat2 = getattr(mat, device)()
+    assert mat2.shape == mat.shape
+    assert torch.allclose(mat2.row, target_row)
+    assert torch.allclose(mat2.col, target_col)
+    assert torch.allclose(mat2.val, target_val)
+
+
+@pytest.mark.parametrize(
+    "dtype", [torch.float, torch.double, torch.int, torch.long]
+)
+def test_to_dtype(dtype):
+    row = torch.tensor([1, 1, 2])
+    col = torch.tensor([1, 2, 0])
+    mat = from_coo(row, col, shape=(3, 4))
+
+    target_val = mat.val.to(dtype=dtype)
+
+    mat2 = mat.to(dtype=dtype)
+    assert mat2.shape == mat.shape
+    assert torch.allclose(mat2.val, target_val)
+
+    func_name = {
+        torch.float: "float",
+        torch.double: "double",
+        torch.int: "int",
+        torch.long: "long",
+    }
+    mat2 = getattr(mat, func_name[dtype])()
+    assert mat2.shape == mat.shape
+    assert torch.allclose(mat2.val, target_val)
+
+
+@pytest.mark.parametrize("dense_dim", [None, 2])
+@pytest.mark.parametrize("row", [[0, 0, 1, 2], (0, 1, 2, 4)])
+@pytest.mark.parametrize("col", [(0, 1, 2, 2), (1, 3, 3, 4)])
+@pytest.mark.parametrize("extra_shape", [(0, 1), (2, 1)])
+def test_sparse_matrix_transpose(dense_dim, row, col, extra_shape):
+    mat_shape = (max(row) + 1 + extra_shape[0], max(col) + 1 + extra_shape[1])
+    val_shape = (len(row),)
+    if dense_dim is not None:
+        val_shape += (dense_dim,)
+    ctx = F.ctx()
+    val = torch.randn(val_shape).to(ctx)
+    row = torch.tensor(row).to(ctx)
+    col = torch.tensor(col).to(ctx)
+    mat = from_coo(row, col, val, mat_shape).transpose()
+    mat_row, mat_col = mat.coo()
+    mat_val = mat.val
+
+    assert mat.shape == mat_shape[::-1]
+    assert torch.allclose(mat_val, val)
+    assert torch.allclose(mat_row, col)
+    assert torch.allclose(mat_col, row)
+
+
+@pytest.mark.parametrize("row", [[0, 0, 1, 2], (0, 1, 2, 4)])
+@pytest.mark.parametrize("col", [(0, 1, 2, 2), (1, 3, 3, 4)])
+@pytest.mark.parametrize("nz_dim", [None, 2])
+@pytest.mark.parametrize("shape", [(5, 5), (6, 7)])
+def test_torch_sparse_coo_conversion(row, col, nz_dim, shape):
+    dev = F.ctx()
+    row = torch.tensor(row).to(dev)
+    col = torch.tensor(col).to(dev)
+    indices = torch.stack([row, col])
+    torch_sparse_shape = shape
+    val_shape = (row.shape[0],)
+    if nz_dim is not None:
+        torch_sparse_shape += (nz_dim,)
+        val_shape += (nz_dim,)
+    val = torch.randn(val_shape).to(dev)
+    torch_sparse_coo = torch.sparse_coo_tensor(indices, val, torch_sparse_shape)
+    spmat = from_torch_sparse(torch_sparse_coo)
+
+    def _assert_spmat_equal_to_torch_sparse_coo(spmat, torch_sparse_coo):
+        assert torch_sparse_coo.layout == torch.sparse_coo
+        # Use .data_ptr() to check whether indices and values are on the same
+        # memory address
+        assert (
+            spmat.indices().data_ptr() == torch_sparse_coo._indices().data_ptr()
+        )
+        assert spmat.val.data_ptr() == torch_sparse_coo._values().data_ptr()
+        assert spmat.shape == torch_sparse_coo.shape[:2]
+
+    _assert_spmat_equal_to_torch_sparse_coo(spmat, torch_sparse_coo)
+    torch_sparse_coo = to_torch_sparse_coo(spmat)
+    _assert_spmat_equal_to_torch_sparse_coo(spmat, torch_sparse_coo)
+
+
+@pytest.mark.parametrize("indptr", [(0, 0, 1, 4), (0, 1, 2, 4)])
+@pytest.mark.parametrize("indices", [(0, 1, 2, 3), (1, 2, 3, 4)])
+@pytest.mark.parametrize("shape", [(3, 5), (3, 7)])
+def test_torch_sparse_csr_conversion(indptr, indices, shape):
+    dev = F.ctx()
+    indptr = torch.tensor(indptr).to(dev)
+    indices = torch.tensor(indices).to(dev)
+    torch_sparse_shape = shape
+    val_shape = (indices.shape[0],)
+    val = torch.randn(val_shape).to(dev)
+    torch_sparse_csr = _torch_sparse_csr_tensor(
+        indptr, indices, val, torch_sparse_shape
+    )
+    spmat = from_torch_sparse(torch_sparse_csr)
+
+    def _assert_spmat_equal_to_torch_sparse_csr(spmat, torch_sparse_csr):
+        indptr, indices, value_indices = spmat.csr()
+        assert torch_sparse_csr.layout == torch.sparse_csr
+        assert value_indices is None
+        # Use .data_ptr() to check whether indices and values are on the same
+        # memory address
+        assert indptr.data_ptr() == torch_sparse_csr.crow_indices().data_ptr()
+        assert indices.data_ptr() == torch_sparse_csr.col_indices().data_ptr()
+        assert spmat.val.data_ptr() == torch_sparse_csr.values().data_ptr()
+        assert spmat.shape == torch_sparse_csr.shape[:2]
+
+    _assert_spmat_equal_to_torch_sparse_csr(spmat, torch_sparse_csr)
+    torch_sparse_csr = to_torch_sparse_csr(spmat)
+    _assert_spmat_equal_to_torch_sparse_csr(spmat, torch_sparse_csr)
+
+
+@pytest.mark.parametrize("indptr", [(0, 0, 1, 4), (0, 1, 2, 4)])
+@pytest.mark.parametrize("indices", [(0, 1, 2, 3), (1, 2, 3, 4)])
+@pytest.mark.parametrize("shape", [(8, 3), (5, 3)])
+def test_torch_sparse_csc_conversion(indptr, indices, shape):
+    dev = F.ctx()
+    indptr = torch.tensor(indptr).to(dev)
+    indices = torch.tensor(indices).to(dev)
+    torch_sparse_shape = shape
+    val_shape = (indices.shape[0],)
+    val = torch.randn(val_shape).to(dev)
+    torch_sparse_csc = torch.sparse_csc_tensor(
+        indptr, indices, val, torch_sparse_shape
+    )
+    spmat = from_torch_sparse(torch_sparse_csc)
+
+    def _assert_spmat_equal_to_torch_sparse_csc(spmat, torch_sparse_csc):
+        indptr, indices, value_indices = spmat.csc()
+        assert torch_sparse_csc.layout == torch.sparse_csc
+        assert value_indices is None
+        # Use .data_ptr() to check whether indices and values are on the same
+        # memory address
+        assert indptr.data_ptr() == torch_sparse_csc.ccol_indices().data_ptr()
+        assert indices.data_ptr() == torch_sparse_csc.row_indices().data_ptr()
+        assert spmat.val.data_ptr() == torch_sparse_csc.values().data_ptr()
+        assert spmat.shape == torch_sparse_csc.shape[:2]
+
+    _assert_spmat_equal_to_torch_sparse_csc(spmat, torch_sparse_csc)
+    torch_sparse_csc = to_torch_sparse_csc(spmat)
+    _assert_spmat_equal_to_torch_sparse_csc(spmat, torch_sparse_csc)
+
+
+### Diag foramt related tests ###
+
+
+@pytest.mark.parametrize("val_shape", [(3,), (3, 2)])
+@pytest.mark.parametrize("mat_shape", [None, (3, 5), (5, 3)])
+def test_diag(val_shape, mat_shape):
+    ctx = F.ctx()
+    # creation
+    val = torch.randn(val_shape).to(ctx)
+    mat = diag(val, mat_shape)
+
+    # val, shape attributes
+    assert torch.allclose(mat.val, val)
+    if mat_shape is None:
+        mat_shape = (val_shape[0], val_shape[0])
+    assert mat.shape == mat_shape
+
+    val = torch.randn(val_shape).to(ctx)
+
+    # nnz
+    assert mat.nnz == val.shape[0]
+    # dtype
+    assert mat.dtype == val.dtype
+    # device
+    assert mat.device == val.device
+
+    # row, col, val
+    edge_index = torch.arange(len(val)).to(mat.device)
+    row, col = mat.coo()
+    val = mat.val
+    assert torch.allclose(row, edge_index)
+    assert torch.allclose(col, edge_index)
+    assert torch.allclose(val, val)
+
+
+@pytest.mark.parametrize("shape", [(3, 3), (3, 5), (5, 3)])
+@pytest.mark.parametrize("d", [None, 2])
+def test_identity(shape, d):
+    ctx = F.ctx()
+    # creation
+    mat = identity(shape, d)
+    # shape
+    assert mat.shape == shape
+    # val
+    len_val = min(shape)
+    if d is None:
+        val_shape = len_val
+    else:
+        val_shape = (len_val, d)
+    val = torch.ones(val_shape)
+    assert torch.allclose(val, mat.val)
+
+
+@pytest.mark.parametrize("val_shape", [(3,), (3, 2)])
+@pytest.mark.parametrize("mat_shape", [None, (3, 5), (5, 3)])
+def test_diag_matrix_transpose(val_shape, mat_shape):
+    ctx = F.ctx()
+    val = torch.randn(val_shape).to(ctx)
+    mat = diag(val, mat_shape).transpose()
+
+    assert torch.allclose(mat.val, val)
+    if mat_shape is None:
+        mat_shape = (val_shape[0], val_shape[0])
+    assert mat.shape == mat_shape[::-1]

tests/python/pytorch/sparse/test_unary_op.py

@@ -0,0 +1,40 @@
+import sys
+
+import backend as F
+import torch
+
+from dgl.sparse import diag, spmatrix
+
+
+def test_neg():
+    ctx = F.ctx()
+    row = torch.tensor([1, 1, 3]).to(ctx)
+    col = torch.tensor([1, 2, 3]).to(ctx)
+    val = torch.tensor([1.0, 1.0, 2.0]).to(ctx)
+    A = spmatrix(torch.stack([row, col]), val)
+    neg_A = -A
+    assert A.shape == neg_A.shape
+    assert A.nnz == neg_A.nnz
+    assert torch.allclose(-A.val, neg_A.val)
+    assert torch.allclose(torch.stack(A.coo()), torch.stack(neg_A.coo()))
+    assert A.val.device == neg_A.val.device
+
+
+def test_diag_neg():
+    ctx = F.ctx()
+    val = torch.arange(3).float().to(ctx)
+    D = diag(val)
+    neg_D = -D
+    assert D.shape == neg_D.shape
+    assert torch.allclose(-D.val, neg_D.val)
+    assert D.val.device == neg_D.val.device
+
+
+def test_diag_inv():
+    ctx = F.ctx()
+    val = torch.arange(1, 4).float().to(ctx)
+    D = diag(val)
+    inv_D = D.inv()
+    assert D.shape == inv_D.shape
+    assert torch.allclose(1.0 / D.val, inv_D.val)
+    assert D.val.device == inv_D.val.device

tests/python/pytorch/sparse/utils.py

@@ -0,0 +1,163 @@
+import numpy as np
+import torch
+
+from dgl.sparse import diag, from_csc, from_csr, SparseMatrix, spmatrix
+
+np.random.seed(42)
+torch.random.manual_seed(42)
+
+
+def clone_detach_and_grad(t):
+    t = t.clone().detach()
+    t.requires_grad_()
+    return t
+
+
+def rand_stride(t):
+    """Add stride to the last dimension of a tensor."""
+    stride = np.random.randint(2, 4)
+    ret = torch.stack([t] * stride, dim=-1)[..., 0]
+    ret = ret.detach()
+    if torch.is_floating_point(t):
+        ret.requires_grad_()
+    return ret
+
+
+def rand_coo(shape, nnz, dev, nz_dim=None):
+    # Create a sparse matrix without duplicate entries.
+    nnzid = np.random.choice(shape[0] * shape[1], nnz, replace=False)
+    nnzid = torch.tensor(nnzid, device=dev).long()
+    row = torch.div(nnzid, shape[1], rounding_mode="floor")
+    col = nnzid % shape[1]
+    if nz_dim is None:
+        val = torch.randn(nnz, device=dev, requires_grad=True)
+    else:
+        val = torch.randn(nnz, nz_dim, device=dev, requires_grad=True)
+    indices = torch.stack([row, col])
+    indices = rand_stride(indices)
+    val = rand_stride(val)
+    return spmatrix(indices, val, shape)
+
+
+def rand_csr(shape, nnz, dev, nz_dim=None):
+    # Create a sparse matrix without duplicate entries.
+    nnzid = np.random.choice(shape[0] * shape[1], nnz, replace=False)
+    nnzid = torch.tensor(nnzid, device=dev).long()
+    row = torch.div(nnzid, shape[1], rounding_mode="floor")
+    col = nnzid % shape[1]
+    if nz_dim is None:
+        val = torch.randn(nnz, device=dev, requires_grad=True)
+    else:
+        val = torch.randn(nnz, nz_dim, device=dev, requires_grad=True)
+    indptr = torch.zeros(shape[0] + 1, device=dev, dtype=torch.int64)
+    for r in row.tolist():
+        indptr[r + 1] += 1
+    indptr = torch.cumsum(indptr, 0)
+    row_sorted, row_sorted_idx = torch.sort(row)
+    indices = col[row_sorted_idx]
+    indptr = rand_stride(indptr)
+    indices = rand_stride(indices)
+    val = rand_stride(val)
+    return from_csr(indptr, indices, val, shape=shape)
+
+
+def rand_csc(shape, nnz, dev, nz_dim=None):
+    # Create a sparse matrix without duplicate entries.
+    nnzid = np.random.choice(shape[0] * shape[1], nnz, replace=False)
+    nnzid = torch.tensor(nnzid, device=dev).long()
+    row = torch.div(nnzid, shape[1], rounding_mode="floor")
+    col = nnzid % shape[1]
+    if nz_dim is None:
+        val = torch.randn(nnz, device=dev, requires_grad=True)
+    else:
+        val = torch.randn(nnz, nz_dim, device=dev, requires_grad=True)
+    indptr = torch.zeros(shape[1] + 1, device=dev, dtype=torch.int64)
+    for c in col.tolist():
+        indptr[c + 1] += 1
+    indptr = torch.cumsum(indptr, 0)
+    col_sorted, col_sorted_idx = torch.sort(col)
+    indices = row[col_sorted_idx]
+    indptr = rand_stride(indptr)
+    indices = rand_stride(indices)
+    val = rand_stride(val)
+    return from_csc(indptr, indices, val, shape=shape)
+
+
+def rand_diag(shape, nnz, dev, nz_dim=None):
+    nnz = min(shape)
+    if nz_dim is None:
+        val = torch.randn(nnz, device=dev, requires_grad=True)
+    else:
+        val = torch.randn(nnz, nz_dim, device=dev, requires_grad=True)
+    return diag(val, shape)
+
+
+def rand_coo_uncoalesced(shape, nnz, dev):
+    # Create a sparse matrix with possible duplicate entries.
+    row = torch.randint(shape[0], (nnz,), device=dev)
+    col = torch.randint(shape[1], (nnz,), device=dev)
+    val = torch.randn(nnz, device=dev, requires_grad=True)
+    indices = torch.stack([row, col])
+    indices = rand_stride(indices)
+    return spmatrix(indices, val, shape)
+
+
+def rand_csr_uncoalesced(shape, nnz, dev):
+    # Create a sparse matrix with possible duplicate entries.
+    row = torch.randint(shape[0], (nnz,), device=dev)
+    col = torch.randint(shape[1], (nnz,), device=dev)
+    val = torch.randn(nnz, device=dev, requires_grad=True)
+    indptr = torch.zeros(shape[0] + 1, device=dev, dtype=torch.int64)
+    for r in row.tolist():
+        indptr[r + 1] += 1
+    indptr = torch.cumsum(indptr, 0)
+    row_sorted, row_sorted_idx = torch.sort(row)
+    indices = col[row_sorted_idx]
+    indptr = rand_stride(indptr)
+    indices = rand_stride(indices)
+    val = rand_stride(val)
+    return from_csr(indptr, indices, val, shape=shape)
+
+
+def rand_csc_uncoalesced(shape, nnz, dev):
+    # Create a sparse matrix with possible duplicate entries.
+    row = torch.randint(shape[0], (nnz,), device=dev)
+    col = torch.randint(shape[1], (nnz,), device=dev)
+    val = torch.randn(nnz, device=dev, requires_grad=True)
+    indptr = torch.zeros(shape[1] + 1, device=dev, dtype=torch.int64)
+    for c in col.tolist():
+        indptr[c + 1] += 1
+    indptr = torch.cumsum(indptr, 0)
+    col_sorted, col_sorted_idx = torch.sort(col)
+    indices = row[col_sorted_idx]
+    indptr = rand_stride(indptr)
+    indices = rand_stride(indices)
+    val = rand_stride(val)
+    return from_csc(indptr, indices, val, shape=shape)
+
+
+def sparse_matrix_to_dense(A: SparseMatrix):
+    dense = A.to_dense()
+    return clone_detach_and_grad(dense)
+
+
+def sparse_matrix_to_torch_sparse(A: SparseMatrix, val=None):
+    row, col = A.coo()
+    edge_index = torch.cat((row.unsqueeze(0), col.unsqueeze(0)), 0)
+    shape = A.shape
+    if val is None:
+        val = A.val
+    val = val.clone().detach()
+    if len(A.val.shape) > 1:
+        shape += (A.val.shape[-1],)
+    ret = torch.sparse_coo_tensor(edge_index, val, shape).coalesce()
+    ret.requires_grad_()
+    return ret
+
+
+def dense_mask(dense, sparse):
+    ret = torch.zeros_like(dense)
+    row, col = sparse.coo()
+    for r, c in zip(row, col):
+        ret[r, c] = dense[r, c]
+    return ret