test_basis.py 1.67 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
2
3
4
from itertools import product

import pytest
import torch
rusty1s's avatar
rusty1s committed
5
from torch_spline_conv import spline_basis
rusty1s's avatar
update  
rusty1s committed
6
from torch_spline_conv.testing import devices, dtypes, tensor
rusty1s's avatar
rusty1s committed
7

rusty1s's avatar
rusty1s committed
8
tests = [{
rusty1s's avatar
rusty1s committed
9
    'pseudo': [[0], [0.0625], [0.25], [0.75], [0.9375], [1]],
rusty1s's avatar
rusty1s committed
10
11
12
13
14
    'kernel_size': [5],
    'is_open_spline': [1],
    'basis': [[1, 0], [0.75, 0.25], [1, 0], [1, 0], [0.25, 0.75], [1, 0]],
    'weight_index': [[0, 1], [0, 1], [1, 2], [3, 4], [3, 4], [4, 0]],
}, {
rusty1s's avatar
rusty1s committed
15
    'pseudo': [[0], [0.0625], [0.25], [0.75], [0.9375], [1]],
rusty1s's avatar
rusty1s committed
16
17
18
19
20
21
22
23
24
25
26
    'kernel_size': [4],
    'is_open_spline': [0],
    'basis': [[1, 0], [0.75, 0.25], [1, 0], [1, 0], [0.25, 0.75], [1, 0]],
    'weight_index': [[0, 1], [0, 1], [1, 2], [3, 0], [3, 0], [0, 1]],
}, {
    'pseudo': [[0.125, 0.5], [0.5, 0.5], [0.75, 0.125]],
    'kernel_size': [5, 5],
    'is_open_spline': [1, 1],
    'basis': [[0.5, 0.5, 0, 0], [1, 0, 0, 0], [0.5, 0, 0.5, 0]],
    'weight_index': [[10, 11, 15, 16], [12, 13, 17, 18], [3, 4, 8, 9]]
}]
rusty1s's avatar
rusty1s committed
27
28


rusty1s's avatar
rusty1s committed
29
30
@pytest.mark.parametrize('test,dtype,device', product(tests, dtypes, devices))
def test_spline_basis_forward(test, dtype, device):
rusty1s's avatar
update  
rusty1s committed
31
32
33
    if dtype == torch.bfloat16 and device == torch.device('cuda:0'):
        return

rusty1s's avatar
rusty1s committed
34
35
36
    pseudo = tensor(test['pseudo'], dtype, device)
    kernel_size = tensor(test['kernel_size'], torch.long, device)
    is_open_spline = tensor(test['is_open_spline'], torch.uint8, device)
rusty1s's avatar
update  
rusty1s committed
37
38
39
    basis = tensor(test['basis'], dtype, device)
    weight_index = tensor(test['weight_index'], dtype, device)

rusty1s's avatar
rusty1s committed
40
    degree = 1
rusty1s's avatar
rusty1s committed
41

rusty1s's avatar
rusty1s committed
42
43
    basis, weight_index = spline_basis(pseudo, kernel_size, is_open_spline,
                                       degree)
rusty1s's avatar
update  
rusty1s committed
44
45
    assert torch.allclose(basis, basis)
    assert torch.allclose(weight_index, weight_index)