test_segment.py 6.27 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.autograd import gradcheck
rusty1s's avatar
rusty1s committed
6
from torch_scatter import segment_csr
rusty1s's avatar
rusty1s committed
7

rusty1s's avatar
rusty1s committed
8
from .utils import tensor, dtypes, devices
rusty1s's avatar
rusty1s committed
9

rusty1s's avatar
rusty1s committed
10
11
reductions = ['sum', 'mean', 'min', 'max']
grad_reductions = ['sum', 'mean']
rusty1s's avatar
rusty1s committed
12
13
14
15
16
17

tests = [
    {
        'src': [1, 2, 3, 4, 5, 6],
        'index': [0, 0, 1, 1, 1, 3],
        'indptr': [0, 2, 5, 5, 6],
rusty1s's avatar
rusty1s committed
18
        'sum': [3, 12, 0, 6],
rusty1s's avatar
rusty1s committed
19
20
21
22
23
24
25
26
27
28
        'mean': [1.5, 4, 0, 6],
        'min': [1, 3, 0, 6],
        'arg_min': [0, 2, 6, 5],
        'max': [2, 5, 0, 6],
        'arg_max': [1, 4, 6, 5],
    },
    {
        'src': [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12]],
        'index': [0, 0, 1, 1, 1, 3],
        'indptr': [0, 2, 5, 5, 6],
rusty1s's avatar
rusty1s committed
29
        'sum': [[4, 6], [21, 24], [0, 0], [11, 12]],
rusty1s's avatar
rusty1s committed
30
31
32
33
34
35
36
37
38
39
        'mean': [[2, 3], [7, 8], [0, 0], [11, 12]],
        'min': [[1, 2], [5, 6], [0, 0], [11, 12]],
        'arg_min': [[0, 0], [2, 2], [6, 6], [5, 5]],
        'max': [[3, 4], [9, 10], [0, 0], [11, 12]],
        'arg_max': [[1, 1], [4, 4], [6, 6], [5, 5]],
    },
    {
        'src': [[1, 3, 5, 7, 9, 11], [2, 4, 6, 8, 10, 12]],
        'index': [[0, 0, 1, 1, 1, 3], [0, 0, 0, 1, 1, 2]],
        'indptr': [[0, 2, 5, 5, 6], [0, 3, 5, 6, 6]],
rusty1s's avatar
rusty1s committed
40
        'sum': [[4, 21, 0, 11], [12, 18, 12, 0]],
rusty1s's avatar
rusty1s committed
41
42
43
44
45
46
47
        'mean': [[2, 7, 0, 11], [4, 9, 12, 0]],
        'min': [[1, 5, 0, 11], [2, 8, 12, 0]],
        'arg_min': [[0, 2, 6, 5], [0, 3, 5, 6]],
        'max': [[3, 9, 0, 11], [6, 10, 12, 0]],
        'arg_max': [[1, 4, 6, 5], [2, 4, 5, 6]],
    },
    {
rusty1s's avatar
rusty1s committed
48
49
50
        'src': [[[1, 2], [3, 4], [5, 6]], [[7, 9], [10, 11], [12, 13]]],
        'index': [[0, 0, 1], [0, 2, 2]],
        'indptr': [[0, 2, 3, 3], [0, 1, 1, 3]],
rusty1s's avatar
rusty1s committed
51
        'sum': [[[4, 6], [5, 6], [0, 0]], [[7, 9], [0, 0], [22, 24]]],
rusty1s's avatar
rusty1s committed
52
53
54
55
56
        'mean': [[[2, 3], [5, 6], [0, 0]], [[7, 9], [0, 0], [11, 12]]],
        'min': [[[1, 2], [5, 6], [0, 0]], [[7, 9], [0, 0], [10, 11]]],
        'arg_min': [[[0, 0], [2, 2], [3, 3]], [[0, 0], [3, 3], [1, 1]]],
        'max': [[[3, 4], [5, 6], [0, 0]], [[7, 9], [0, 0], [12, 13]]],
        'arg_max': [[[1, 1], [2, 2], [3, 3]], [[0, 0], [3, 3], [2, 2]]],
rusty1s's avatar
rusty1s committed
57
58
59
60
61
    },
    {
        'src': [[1, 3], [2, 4]],
        'index': [[0, 0], [0, 0]],
        'indptr': [[0, 2], [0, 2]],
rusty1s's avatar
rusty1s committed
62
        'sum': [[4], [6]],
rusty1s's avatar
rusty1s committed
63
64
65
66
67
68
69
70
71
72
        'mean': [[2], [3]],
        'min': [[1], [2]],
        'arg_min': [[0], [0]],
        'max': [[3], [4]],
        'arg_max': [[1], [1]],
    },
    {
        'src': [[[1, 1], [3, 3]], [[2, 2], [4, 4]]],
        'index': [[0, 0], [0, 0]],
        'indptr': [[0, 2], [0, 2]],
rusty1s's avatar
rusty1s committed
73
        'sum': [[[4, 4]], [[6, 6]]],
rusty1s's avatar
rusty1s committed
74
75
76
77
78
79
80
81
82
83
84
        'mean': [[[2, 2]], [[3, 3]]],
        'min': [[[1, 1]], [[2, 2]]],
        'arg_min': [[[0, 0]], [[0, 0]]],
        'max': [[[3, 3]], [[4, 4]]],
        'arg_max': [[[1, 1]], [[1, 1]]],
    },
]


@pytest.mark.parametrize('test,reduce,dtype,device',
                         product(tests, reductions, dtypes, devices))
rusty1s's avatar
rusty1s committed
85
def test_forward(test, reduce, dtype, device):
rusty1s's avatar
rusty1s committed
86
87
88
89
90
    src = tensor(test['src'], dtype, device)
    index = tensor(test['index'], torch.long, device)
    indptr = tensor(test['indptr'], torch.long, device)
    expected = tensor(test[reduce], dtype, device)

rusty1s's avatar
rusty1s committed
91
92
93
94
95
96
    # out = segment_coo(src, index, reduce=reduce)
    # if isinstance(out, tuple):
    #     out, arg_out = out
    #     arg_expected = tensor(test[f'arg_{reduce}'], torch.long, device)
    #     assert torch.all(arg_out == arg_expected)
    # assert torch.all(out == expected)
rusty1s's avatar
rusty1s committed
97
98
99
100
101
102
103
104

    out = segment_csr(src, indptr, reduce=reduce)
    if isinstance(out, tuple):
        out, arg_out = out
        arg_expected = tensor(test[f'arg_{reduce}'], torch.long, device)
        assert torch.all(arg_out == arg_expected)
    assert torch.all(out == expected)

rusty1s's avatar
rusty1s committed
105

rusty1s's avatar
rusty1s committed
106
107
108
109
110
111
112
113
@pytest.mark.parametrize('test,reduce,device',
                         product(tests, grad_reductions, devices))
def test_backward(test, reduce, device):
    src = tensor(test['src'], torch.double, device)
    src.requires_grad_()
    index = tensor(test['index'], torch.long, device)
    indptr = tensor(test['indptr'], torch.long, device)

rusty1s's avatar
rusty1s committed
114
    # assert gradcheck(segment_coo, (src, index, None, None, reduce))
rusty1s's avatar
typo  
rusty1s committed
115
    assert gradcheck(segment_csr, (src, indptr, None, reduce))
rusty1s's avatar
rusty1s committed
116
117


rusty1s's avatar
rusty1s committed
118
119
120
121
122
123
124
125
126
127
128
129
@pytest.mark.parametrize('test,reduce,dtype,device',
                         product(tests, reductions, dtypes, devices))
def test_segment_out(test, reduce, dtype, device):
    src = tensor(test['src'], dtype, device)
    index = tensor(test['index'], torch.long, device)
    indptr = tensor(test['indptr'], torch.long, device)
    expected = tensor(test[reduce], dtype, device)

    size = list(src.size())
    size[indptr.dim() - 1] = indptr.size(-1) - 1
    out = src.new_full(size, -2)

rusty1s's avatar
rusty1s committed
130
    segment_csr(src, indptr, out, reduce=reduce)
rusty1s's avatar
rusty1s committed
131
132
    assert torch.all(out == expected)

rusty1s's avatar
rusty1s committed
133
    # out.fill_(-2)
rusty1s's avatar
rusty1s committed
134

rusty1s's avatar
rusty1s committed
135
    # segment_coo(src, index, out, reduce=reduce)
rusty1s's avatar
rusty1s committed
136

rusty1s's avatar
rusty1s committed
137
138
139
140
141
142
143
144
145
146
    # if reduce == 'sum':
    #     expected = expected - 2
    # elif reduce == 'mean':
    #     expected = out  # We can not really test this here.
    # elif reduce == 'min':
    #     expected = expected.fill_(-2)
    # elif reduce == 'max':
    #     expected[expected == 0] = -2
    # else:
    #     raise ValueError
rusty1s's avatar
rusty1s committed
147

rusty1s's avatar
rusty1s committed
148
    # assert torch.all(out == expected)
rusty1s's avatar
rusty1s committed
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165


@pytest.mark.parametrize('test,reduce,dtype,device',
                         product(tests, reductions, dtypes, devices))
def test_non_contiguous_segment(test, reduce, dtype, device):
    src = tensor(test['src'], dtype, device)
    index = tensor(test['index'], torch.long, device)
    indptr = tensor(test['indptr'], torch.long, device)
    expected = tensor(test[reduce], dtype, device)

    if src.dim() > 1:
        src = src.transpose(0, 1).contiguous().transpose(0, 1)
    if index.dim() > 1:
        index = index.transpose(0, 1).contiguous().transpose(0, 1)
    if indptr.dim() > 1:
        indptr = indptr.transpose(0, 1).contiguous().transpose(0, 1)

rusty1s's avatar
rusty1s committed
166
167
168
169
170
171
    # out = segment_coo(src, index, reduce=reduce)
    # if isinstance(out, tuple):
    #     out, arg_out = out
    #     arg_expected = tensor(test[f'arg_{reduce}'], torch.long, device)
    #     assert torch.all(arg_out == arg_expected)
    # assert torch.all(out == expected)
rusty1s's avatar
rusty1s committed
172
173
174
175
176
177
178

    out = segment_csr(src, indptr, reduce=reduce)
    if isinstance(out, tuple):
        out, arg_out = out
        arg_expected = tensor(test[f'arg_{reduce}'], torch.long, device)
        assert torch.all(arg_out == arg_expected)
    assert torch.all(out == expected)