test_std.py 1005 Bytes
Newer Older
rusty1s's avatar
typo  
rusty1s committed
1
2
3
4
5
6
7
8
from itertools import product

import pytest
import torch
from torch_scatter import scatter_std

from .utils import grad_dtypes as dtypes, devices, tensor

rusty1s's avatar
rusty1s committed
9
biases = [True, False]
rusty1s's avatar
typo  
rusty1s committed
10
11


rusty1s's avatar
rusty1s committed
12
@pytest.mark.parametrize('dtype,device,bias', product(dtypes, devices, biases))
rusty1s's avatar
typo  
rusty1s committed
13
14
15
16
17
def test_std(dtype, device, bias):
    src = tensor([[2, 0, 1, 4, 3], [0, 2, 1, 3, 4]], dtype, device)
    index = tensor([[0, 0, 0, 0, 0], [1, 1, 1, 1, 1]], torch.long, device)

    out = scatter_std(src, index, dim=-1, unbiased=bias)
rusty1s's avatar
rusty1s committed
18
    std = src.std(dim=-1, unbiased=bias)[0].item()
19
    expected = tensor([[std, 0], [0, std]], dtype, device)
rusty1s's avatar
rusty1s committed
20
    assert torch.allclose(out, expected)
rusty1s's avatar
rusty1s committed
21
22
23
24


@pytest.mark.parametrize('dtype,device', product(dtypes, devices))
def test_empty_std(dtype, device):
25
26
27
    out = torch.zeros(1, 5, dtype=dtype, device=device)
    src = tensor([], dtype, device).view(0, 5)
    index = tensor([], torch.long, device).view(0, 5)
rusty1s's avatar
rusty1s committed
28

29
30
    out = scatter_std(src, index, dim=0, out=out)
    assert out.tolist() == [[0, 0, 0, 0, 0]]