test_unary_op_sp.py 520 Bytes
Newer Older
Mufei Li's avatar
Mufei Li committed
1
2
3
import sys

import backend as F
4
import torch
Mufei Li's avatar
Mufei Li committed
5

6
from dgl.sparse import from_coo
Mufei Li's avatar
Mufei Li committed
7
8
9
10
11
12
13


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)
14
    A = from_coo(row, col, val)
Mufei Li's avatar
Mufei Li committed
15
16
17
18
19
20
    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