fix test (#340)

test/test_broadcasting.py

@@ -3,8 +3,7 @@ from itertools import product
 import pytest
 import torch
 from torch_scatter import scatter
-
-from .utils import reductions, devices
+from torch_scatter.testing import devices, reductions
 
 
 @pytest.mark.parametrize('reduce,device', product(reductions, devices))

test/test_gather.py

@@ -3,9 +3,8 @@ from itertools import product
 import pytest
 import torch
 from torch.autograd import gradcheck
-from torch_scatter import gather_csr, gather_coo
-
-from .utils import tensor, dtypes, devices
+from torch_scatter import gather_coo, gather_csr
+from torch_scatter.testing import devices, dtypes, tensor
 
 tests = [
     {

test/test_multi_gpu.py

@@ -3,8 +3,7 @@ from itertools import product
 import pytest
 import torch
 import torch_scatter
-
-from .utils import reductions, tensor, dtypes
+from torch_scatter.testing import dtypes, reductions, tensor
 
 tests = [
     {

test/test_scatter.py

@@ -4,8 +4,7 @@ import pytest
 import torch
 import torch_scatter
 from torch.autograd import gradcheck
-
-from .utils import devices, dtypes, reductions, tensor
+from torch_scatter.testing import devices, dtypes, reductions, tensor
 
 reductions = reductions + ['mul']
 

test/test_segment.py

@@ -2,10 +2,9 @@ from itertools import product
 
 import pytest
 import torch
-from torch.autograd import gradcheck
 import torch_scatter
-
-from .utils import reductions, tensor, dtypes, devices
+from torch.autograd import gradcheck
+from torch_scatter.testing import devices, dtypes, reductions, tensor
 
 tests = [
     {

test/test_zero_tensors.py

@@ -2,10 +2,9 @@ from itertools import product
 
 import pytest
 import torch
-from torch_scatter import scatter, segment_coo, gather_coo
-from torch_scatter import segment_csr, gather_csr
-
-from .utils import reductions, tensor, grad_dtypes, devices
+from torch_scatter import (gather_coo, gather_csr, scatter, segment_coo,
+                           segment_csr)
+from torch_scatter.testing import devices, grad_dtypes, reductions, tensor
 
 
 @pytest.mark.parametrize('reduce,dtype,device',

test/utils.py → torch_scatter/testing.py

+from typing import Any
+
 import torch
 
 reductions = ['sum', 'add', 'mean', 'min', 'max']
 
-dtypes = [torch.half, torch.bfloat16, torch.float, torch.double,
-          torch.int, torch.long]
+dtypes = [
+    torch.half, torch.bfloat16, torch.float, torch.double, torch.int,
+    torch.long
+]
 grad_dtypes = [torch.float, torch.double]
 
 devices = [torch.device('cpu')]
 if torch.cuda.is_available():
-    devices += [torch.device(f'cuda:{torch.cuda.current_device()}')]
+    devices += [torch.device('cuda:0')]
 
 
-def tensor(x, dtype, device):
+def tensor(x: Any, dtype: torch.dtype, device: torch.device):
     return None if x is None else torch.tensor(x, device=device).to(dtype)