Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
test_consecutive.py 993 Bytes
Newer Older
rusty1s's avatar
cuda tests  
rusty1s committed
1 
import pytest
rusty1s's avatar
tests  
rusty1s committed
2
3
4
5
6
7
8
9
10
11
12
13 
import torch
from torch_cluster.functions.utils.consecutive import consecutive


def test_consecutive():
    vec = torch.LongTensor([0, 2, 3])
    assert consecutive(vec).tolist() == [0, 1, 2]

    vec = torch.LongTensor([0, 3, 2, 2, 3])
    assert consecutive(vec).tolist() == [0, 2, 1, 1, 2]

    vec = torch.LongTensor([0, 3, 2, 2, 3])
rusty1s's avatar
cuda tests  
rusty1s committed
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28 
    assert consecutive(vec, True)[0].tolist() == [0, 2, 1, 1, 2]
    assert consecutive(vec, True)[1].tolist() == [0, 2, 3]


@pytest.mark.skipif(not torch.cuda.is_available(), reason='no CUDA')
def test_consecutive_gpu():  # pragma: no cover
    vec = torch.cuda.LongTensor([0, 2, 3])
    assert consecutive(vec).cpu().tolist() == [0, 1, 2]

    vec = torch.cuda.LongTensor([0, 3, 2, 2, 3])
    assert consecutive(vec).cpu().tolist() == [0, 2, 1, 1, 2]

    vec = torch.cuda.LongTensor([0, 3, 2, 2, 3])
    assert consecutive(vec, True)[0].cpu().tolist() == [0, 2, 1, 1, 2]
    assert consecutive(vec, True)[1].cpu().tolist() == [0, 2, 3]