update code and tests

Author: rusty1s

.coveragerc

@@ -3,10 +3,5 @@ source=torch_scatter
 [report]
 exclude_lines =
     pragma: no cover
-    cuda
-    forward
-    backward
-    apply
+    torch.jit.script
     raise
-    min_value
-    max_value

test/composite/test_logsumexp.py

@@ -0,0 +1,18 @@
+import torch
+from torch_scatter import scatter_logsumexp
+
+
+def test_logsumexp():
+    src = torch.tensor([0.5, 0, 0.5, -2.1, 3.2, 7, -1, -100])
+    index = torch.tensor([0, 1, 0, 1, 1, 2, 4, 4])
+
+    out = scatter_logsumexp(src, index)
+
+    out0 = torch.logsumexp(torch.tensor([0.5, 0.5]), dim=-1)
+    out1 = torch.logsumexp(torch.tensor([0, -2.1, 3.2]), dim=-1)
+    out2 = torch.logsumexp(torch.tensor(7, dtype=torch.float), dim=-1)
+    out3 = torch.logsumexp(torch.tensor([], dtype=torch.float), dim=-1)
+    out4 = torch.tensor(-1, dtype=torch.float)
+
+    expected = torch.stack([out0, out1, out2, out3, out4], dim=0)
+    assert torch.allclose(out, expected)

test/composite/test_softmax.py

@@ -1,57 +1,38 @@
-from itertools import product
-
-import pytest
 import torch
-from torch_scatter.composite import scatter_log_softmax, scatter_softmax
-
-from test.utils import devices, tensor, grad_dtypes
+from torch_scatter import scatter_log_softmax, scatter_softmax
 
 
-@pytest.mark.parametrize('dtype,device', product(grad_dtypes, devices))
-def test_softmax(dtype, device):
-    src = tensor([0.2, 0, 0.2, -2.1, 3.2, 7, -1, float('-inf')], dtype, device)
-    index = tensor([0, 1, 0, 1, 1, 2, 4, 4], torch.long, device)
+def test_softmax():
+    src = torch.tensor([0.2, 0, 0.2, -2.1, 3.2, 7, -1, float('-inf')])
+    index = torch.tensor([0, 1, 0, 1, 1, 2, 4, 4])
 
     out = scatter_softmax(src, index)
 
-    out0 = torch.softmax(torch.tensor([0.2, 0.2], dtype=dtype), dim=-1)
-    out1 = torch.softmax(torch.tensor([0, -2.1, 3.2], dtype=dtype), dim=-1)
-    out2 = torch.softmax(torch.tensor([7], dtype=dtype), dim=-1)
-    out4 = torch.softmax(torch.tensor([-1, float('-inf')], dtype=dtype),
-                         dim=-1)
+    out0 = torch.softmax(torch.tensor([0.2, 0.2]), dim=-1)
+    out1 = torch.softmax(torch.tensor([0, -2.1, 3.2]), dim=-1)
+    out2 = torch.softmax(torch.tensor([7], dtype=torch.float), dim=-1)
+    out4 = torch.softmax(torch.tensor([-1, float('-inf')]), dim=-1)
 
     expected = torch.stack([
         out0[0], out1[0], out0[1], out1[1], out1[2], out2[0], out4[0], out4[1]
-    ], dim=0).to(device)
+    ], dim=0)
 
     assert torch.allclose(out, expected)
 
 
-@pytest.mark.parametrize('dtype,device', product(grad_dtypes, devices))
-def test_softmax_broadcasting(dtype, device):
-    src = torch.randn(10, 5, dtype=dtype, device=device)
-    index = tensor([0, 0, 1, 1, 2, 2, 3, 3, 4, 4], torch.long, device)
-
-    out = scatter_softmax(src, index, dim=0).view(5, 2, 5)
-    out = out.sum(dim=1)
-    assert torch.allclose(out, torch.ones_like(out))
-
-
-@pytest.mark.parametrize('dtype,device', product(grad_dtypes, devices))
-def test_log_softmax(dtype, device):
-    src = tensor([0.2, 0, 0.2, -2.1, 3.2, 7, -1, float('-inf')], dtype, device)
-    index = tensor([0, 1, 0, 1, 1, 2, 4, 4], torch.long, device)
+def test_log_softmax():
+    src = torch.tensor([0.2, 0, 0.2, -2.1, 3.2, 7, -1, float('-inf')])
+    index = torch.tensor([0, 1, 0, 1, 1, 2, 4, 4])
 
     out = scatter_log_softmax(src, index)
 
-    out0 = torch.log_softmax(torch.tensor([0.2, 0.2], dtype=dtype), dim=-1)
-    out1 = torch.log_softmax(torch.tensor([0, -2.1, 3.2], dtype=dtype), dim=-1)
-    out2 = torch.log_softmax(torch.tensor([7], dtype=dtype), dim=-1)
-    out4 = torch.log_softmax(torch.tensor([-1, float('-inf')], dtype=dtype),
-                             dim=-1)
+    out0 = torch.log_softmax(torch.tensor([0.2, 0.2]), dim=-1)
+    out1 = torch.log_softmax(torch.tensor([0, -2.1, 3.2]), dim=-1)
+    out2 = torch.log_softmax(torch.tensor([7], dtype=torch.float), dim=-1)
+    out4 = torch.log_softmax(torch.tensor([-1, float('-inf')]), dim=-1)
 
     expected = torch.stack([
         out0[0], out1[0], out0[1], out1[1], out1[2], out2[0], out4[0], out4[1]
-    ], dim=0).to(device)
+    ], dim=0)
 
     assert torch.allclose(out, expected)

test/composite/test_std.py

@@ -0,0 +1,12 @@
+import torch
+from torch_scatter import scatter_std
+
+
+def test_std():
+    src = torch.tensor([[2, 0, 1, 4, 3], [0, 2, 1, 3, 4]], dtype=torch.float)
+    index = torch.tensor([[0, 0, 0, 0, 0], [1, 1, 1, 1, 1]], dtype=torch.long)
+
+    out = scatter_std(src, index, dim=-1, unbiased=True)
+    std = src.std(dim=-1, unbiased=True)[0]
+    expected = torch.tensor([[std, 0], [0, std]])
+    assert torch.allclose(out, expected)

test/test_backward.py

@@ -14,16 +14,6 @@ def test_broadcasting(device):
     out = scatter_add(src, index, dim=2, dim_size=H)
     assert out.size() == (B, C, H, W)
 
-    src = torch.randn((B, 1, H, W), device=device)
-    index = torch.randint(0, H, (B, C, H, W)).to(device, torch.long)
-    out = scatter_add(src, index, dim=2, dim_size=H)
-    assert out.size() == (B, C, H, W)
-
-    src = torch.randn((B, 1, H, W), device=device)
-    index = torch.randint(0, H, (B, 1, H, W)).to(device, torch.long)
-    out = scatter_add(src, index, dim=2, dim_size=H)
-    assert out.size() == (B, 1, H, W)
-
     src = torch.randn((B, C, H, W), device=device)
     index = torch.randint(0, H, (H, )).to(device, torch.long)
     out = scatter_add(src, index, dim=2, dim_size=H)