test_degree.py 1.01 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
import pytest
rusty1s's avatar
tests  
rusty1s committed
2
3
4
5
import torch
from torch_cluster.functions.utils.degree import node_degree


rusty1s's avatar
rusty1s committed
6
def test_node_degree_cpu():
rusty1s's avatar
rusty1s committed
7
8
    index = torch.LongTensor([0, 1, 1, 0, 0, 3, 0])
    degree = node_degree(index, 4)
rusty1s's avatar
tests  
rusty1s committed
9
    expected_degree = [4, 2, 0, 1]
rusty1s's avatar
rusty1s committed
10
    assert degree.type() == torch.LongTensor().type()
rusty1s's avatar
tests  
rusty1s committed
11
12
    assert degree.tolist() == expected_degree

rusty1s's avatar
rusty1s committed
13
    degree = node_degree(index, 4, out=torch.FloatTensor())
rusty1s's avatar
rusty1s committed
14
    assert degree.type() == torch.FloatTensor().type()
rusty1s's avatar
tests  
rusty1s committed
15
    assert degree.tolist() == expected_degree
rusty1s's avatar
rusty1s committed
16
17
18
19


@pytest.mark.skipif(not torch.cuda.is_available(), reason='no CUDA')
def test_node_degree_gpu():  # pragma: no cover
rusty1s's avatar
rusty1s committed
20
21
    index = torch.cuda.LongTensor([0, 1, 1, 0, 0, 3, 0])
    degree = node_degree(index, 4)
rusty1s's avatar
rusty1s committed
22
23
24
25
    expected_degree = [4, 2, 0, 1]
    assert degree.type() == torch.cuda.LongTensor().type()
    assert degree.cpu().tolist() == expected_degree

rusty1s's avatar
rusty1s committed
26
    degree = node_degree(index, 4, out=torch.cuda.FloatTensor())
rusty1s's avatar
rusty1s committed
27
28
    assert degree.type() == torch.cuda.FloatTensor().type()
    assert degree.cpu().tolist() == expected_degree