new pytorch 0.4.0 format

Author: rusty1s

setup.py

@@ -2,7 +2,7 @@
 
 from setuptools import setup, find_packages
 
-__version__ = '0.3.0'
+__version__ = '1.0.0'
 url = 'https://github.com/rusty1s/pytorch_scatter'
 
 install_requires = ['cffi']

test/test_backward.py

@@ -1,49 +1,22 @@
-from os import path as osp
 from itertools import product
 
 import pytest
-import json
 import torch
-from torch.autograd import Variable as V
+from torch.autograd import gradcheck
 import torch_scatter
 
-from .utils import tensors, Tensor
+from .utils import devices
 
-f = open(osp.join(osp.dirname(__file__), 'backward.json'), 'r')
-data = json.load(f)
-f.close()
+funcs = ['add']
+indices = [2, 0, 1, 1, 0]
 
 
-@pytest.mark.parametrize('tensor,i', product(tensors, range(len(data))))
-def test_backward_cpu(tensor, i):
-    name = data[i]['name']
-    index = V(torch.LongTensor(data[i]['index']))
-    input = V(Tensor(tensor, data[i]['input']), requires_grad=True)
-    dim = data[i]['dim']
-    fill_value = data[i]['fill_value']
-    grad = Tensor(tensor, data[i]['grad'])
-    output = V(grad.new(grad.size()).fill_(fill_value))
-    expected = Tensor(tensor, data[i]['expected'])
+@pytest.mark.parametrize('func,device', product(funcs, devices))
+def test_backward(func, device):
+    index = torch.tensor(indices, dtype=torch.long, device=device)
+    src = torch.rand(index.size(), dtype=torch.double, device=device)
+    src.requires_grad_()
 
-    func = getattr(torch_scatter, 'scatter_{}_'.format(name))
-    func(output, index, input, dim)
-    output.backward(grad)
-    assert input.grad.data.tolist() == expected.tolist()
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason='no CUDA')
-@pytest.mark.parametrize('tensor,i', product(tensors, range(len(data))))
-def test_backward_gpu(tensor, i):  # pragma: no cover
-    name = data[i]['name']
-    index = V(torch.cuda.LongTensor(data[i]['index']))
-    input = V(Tensor(tensor, data[i]['input']).cuda(), requires_grad=True)
-    dim = data[i]['dim']
-    fill_value = data[i]['fill_value']
-    grad = Tensor(tensor, data[i]['grad']).cuda()
-    output = V(grad.new(grad.size()).fill_(fill_value).cuda())
-    expected = Tensor(tensor, data[i]['expected'])
-
-    func = getattr(torch_scatter, 'scatter_{}_'.format(name))
-    func(output, index, input, dim)
-    output.backward(grad)
-    assert input.grad.data.cpu().tolist() == expected.tolist()
+    op = getattr(torch_scatter, 'scatter_{}'.format(func))
+    data = (src, index)
+    assert gradcheck(op, data, eps=1e-6, atol=1e-4) is True

test/test_forward.py

@@ -1,67 +1,49 @@
-from os import path as osp
 from itertools import product
 
 import pytest
-import json
 import torch
 import torch_scatter
 
-from .utils import tensors, Tensor
-
-f = open(osp.join(osp.dirname(__file__), 'forward.json'), 'r')
-data = json.load(f)
-f.close()
-
-
-@pytest.mark.parametrize('tensor,i', product(tensors, range(len(data))))
-def test_forward_cpu(tensor, i):
-    name = data[i]['name']
-    index = torch.LongTensor(data[i]['index'])
-    input = Tensor(tensor, data[i]['input'])
-    dim = data[i]['dim']
-    fill_value = data[i]['fill_value']
-    expected = torch.FloatTensor(data[i]['expected']).type_as(input)
-    output = expected.new(expected.size()).fill_(fill_value)
-
-    func = getattr(torch_scatter, 'scatter_{}_'.format(name))
-    result = func(output, index, input, dim)
-    assert output.tolist() == expected.tolist()
-    if 'expected_arg' in data[i]:
-        expected_arg = torch.LongTensor(data[i]['expected_arg'])
-        assert result[1].tolist() == expected_arg.tolist()
-
-    func = getattr(torch_scatter, 'scatter_{}'.format(name))
-    result = func(index, input, dim, fill_value=fill_value)
-    if 'expected_arg' not in data[i]:
-        assert result.tolist() == expected.tolist()
-    else:
-        expected_arg = torch.LongTensor(data[i]['expected_arg'])
-        assert result[0].tolist() == expected.tolist()
-        assert result[1].tolist() == expected_arg.tolist()
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason='no CUDA')
-@pytest.mark.parametrize('tensor,i', product(tensors, range(len(data))))
-def test_forward_gpu(tensor, i):  # pragma: no cover
-    name = data[i]['name']
-    index = torch.cuda.LongTensor(data[i]['index'])
-    input = Tensor(tensor, data[i]['input']).cuda()
-    dim = data[i]['dim']
-    fill_value = data[i]['fill_value']
-    expected = torch.FloatTensor(data[i]['expected']).type_as(input)
-    output = expected.new(expected.size()).fill_(fill_value).cuda()
-
-    func = getattr(torch_scatter, 'scatter_{}_'.format(name))
-    result = func(output, index, input, dim)
-    assert output.cpu().tolist() == expected.tolist()
-    if 'expected_arg' in data[i]:
-        expected_arg = torch.LongTensor(data[i]['expected_arg'])
-        assert result[1].cpu().tolist() == expected_arg.tolist()
-    func = getattr(torch_scatter, 'scatter_{}'.format(name))
-    result = func(index, input, dim, fill_value=fill_value)
-    if 'expected_arg' not in data[i]:
-        assert result.cpu().tolist() == expected.tolist()
-    else:
-        expected_arg = torch.LongTensor(data[i]['expected_arg'])
-        assert result[0].cpu().tolist() == expected.tolist()
-        assert result[1].cpu().tolist() == expected_arg.tolist()
+from .utils import dtypes, devices, tensor
+
+tests = [{
+    'name': 'add',
+    'src': [[2, 0, 1, 4, 3], [0, 2, 1, 3, 4]],
+    'index': [[4, 5, 4, 2, 3], [0, 0, 2, 2, 1]],
+    'fill_value': 0,
+    'expected': [[0, 0, 4, 3, 3, 0], [2, 4, 4, 0, 0, 0]]
+}]
+
+
+@pytest.mark.parametrize('test,dtype,device', product(tests, dtypes, devices))
+def test_forward(test, dtype, device):
+    src = tensor(test['src'], dtype, device)
+    index = tensor(test['index'], torch.long, device)
+
+    op = getattr(torch_scatter, 'scatter_{}'.format(test['name']))
+    output = op(src, index, fill_value=test['fill_value'])
+
+    assert output.tolist() == test['expected']
+    # name = data[i]['name']
+    # index = torch.LongTensor(data[i]['index'])
+    # input = Tensor(tensor, data[i]['input'])
+    # dim = data[i]['dim']
+    # fill_value = data[i]['fill_value']
+    # expected = torch.FloatTensor(data[i]['expected']).type_as(input)
+    # output = expected.new(expected.size()).fill_(fill_value)
+
+    # func = getattr(torch_scatter, 'scatter_{}_'.format(name))
+    # result = func(output, index, input, dim)
+    # assert output.tolist() == expected.tolist()
+    # if 'expected_arg' in data[i]:
+    #     expected_arg = torch.LongTensor(data[i]['expected_arg'])
+    #     assert result[1].tolist() == expected_arg.tolist()
+
+    # func = getattr(torch_scatter, 'scatter_{}'.format(name))
+    # result = func(index, input, dim, fill_value=fill_value)
+    # if 'expected_arg' not in data[i]:
+    #     assert result.tolist() == expected.tolist()
+    # else:
+    #     expected_arg = torch.LongTensor(data[i]['expected_arg'])
+    #     assert result[0].tolist() == expected.tolist()
+    #     assert result[1].tolist() == expected_arg.tolist()

test/utils.py

@@ -1,12 +1,16 @@
 import torch
-from torch._tensor_docs import tensor_classes
+from torch.testing import get_all_dtypes
 
-tensors = [t[:-4] for t in tensor_classes]
-tensors.remove('ShortTensor')  # TODO: PyTorch `atomicAdd` bug with short type.
-tensors.remove('ByteTensor')  # We cannot properly test unsigned values.
-tensors.remove('CharTensor')  # Overflow on gradient computations :(
+dtypes = get_all_dtypes()
+dtypes.remove(torch.half)
+dtypes.remove(torch.short)  # TODO: PyTorch `atomicAdd` bug with short type.
+dtypes.remove(torch.uint8)  # We cannot properly test unsigned values.
+dtypes.remove(torch.int8)  # Overflow on gradient computations :(
 
+devices = [torch.device('cpu')]
+if torch.cuda.is_available():  # pragma: no cover
+    devices += [torch.device('cuda:{}'.format(torch.cuda.current_device()))]
 
-def Tensor(str, x):
-    tensor = getattr(torch, str)
-    return tensor(x)
+
+def tensor(x, dtype, device):
+    return None if x is None else torch.tensor(x, dtype=dtype, device=device)

torch_scatter/__init__.py

@@ -1,16 +1,5 @@
-from .functions.add import scatter_add_, scatter_add
-from .functions.sub import scatter_sub_, scatter_sub
-from .functions.mul import scatter_mul_, scatter_mul
-from .functions.div import scatter_div_, scatter_div
-from .functions.mean import scatter_mean_, scatter_mean
-from .functions.max import scatter_max_, scatter_max
-from .functions.min import scatter_min_, scatter_min
+from .add import ScatterAdd, scatter_add
 
-__version__ = '0.3.0'
+__version__ = '1.0.0'
 
-__all__ = [
-    'scatter_add_', 'scatter_add', 'scatter_sub_', 'scatter_sub',
-    'scatter_mul_', 'scatter_mul', 'scatter_div_', 'scatter_div',
-    'scatter_mean_', 'scatter_mean', 'scatter_max_', 'scatter_max',
-    'scatter_min_', 'scatter_min', '__version__'
-]
+__all__ = ['ScatterAdd', 'scatter_add', '__version__']

torch_scatter/add.py

@@ -0,0 +1,90 @@
+from torch.autograd import Function
+
+from .utils.gen import gen
+
+
+class ScatterAdd(Function):
+    @staticmethod
+    def forward(ctx, out, src, index, dim=-1):
+        ctx.mark_dirty(out)
+        ctx.save_for_backward(index)
+        return out.scatter_add_(dim, index, src)
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        index, = ctx.saved_variables
+
+        grad_src = None
+        if ctx.needs_input_grad[1]:
+            grad_src = grad_out[index]
+
+        return None, grad_src, None, None
+
+
+def scatter_add(src, index, dim=-1, out=None, dim_size=None, fill_value=0):
+    r"""
+    |
+
+    .. image:: https://raw.githubusercontent.com/rusty1s/pytorch_scatter/
+            master/docs/source/_figures/add.svg?sanitize=true
+        :align: center
+        :width: 400px
+
+    |
+
+    Sums all values from the :attr:`src` tensor into :attr:`out` at the indices
+    specified in the :attr:`index` tensor along an given axis :attr:`dim`. For
+    each value in :attr:`src`, its output index is specified by its index in
+    :attr:`input` for dimensions outside of :attr:`dim` and by the
+    corresponding value in :attr:`index` for dimension :attr:`dim`. If
+    multiple indices reference the same location, their **contributions add**.
+
+    Formally, if :attr:`src` and :attr:`index` are n-dimensional tensors with
+    size :math:`(x_0, ..., x_{i-1}, x_i, x_{i+1}, ..., x_{n-1})` and
+    :attr:`dim` = `i`, then :attr:`out` must be an n-dimensional tensor with
+    size :math:`(x_0, ..., x_{i-1}, y, x_{i+1}, ..., x_{n-1})`. Moreover, the
+    values of :attr:`index` must be between `0` and `out.size(dim) - 1`.
+
+    For one-dimensional tensors, the operation computes
+
+    .. math::
+        \mathrm{out}_i = \mathrm{out}_i + \sum_j \mathrm{src}_j
+
+    where sum is over :math:`j` such that :math:`\mathrm{index}_j = i`.
+
+    Args:
+        src (Tensor): The source tensor.
+        index (LongTensor): The indices of elements to scatter.
+        dim (int, optional): The axis along which to index.
+            (default: :obj:`-1`)
+        out (Tensor, optional): The destination tensor. (default: :obj:`None`)
+        dim_size (int, optional): If :attr:`out` is not given, automatically
+            create output with size :attr:`dim_size` at dimension :attr:`dim`.
+            If :attr:`dim_size` is not given, a minimal sized output tensor is
+            returned. (default: :obj:`None`)
+        fill_value (int, optional): If :attr:`out` is not given, automatically
+            fill output tensor with :attr:`fill_value`. (default: :obj:`0`)
+
+    :rtype: :class:`Tensor`
+
+    .. testsetup::
+
+        import torch
+
+    .. testcode::
+
+        from torch_scatter import scatter_add_
+        src = torch.tensor([[2, 0, 1, 4, 3], [0, 2, 1, 3, 4]])
+        index = torch.tensor([[4, 5, 4, 2, 3], [0, 0, 2, 2, 1]])
+        out = src.new_zeros((2, 6))
+        out = scatter_add(src, index, out=out)
+        print(out)
+
+    .. testoutput::
+
+        0  0  4  3  3  0
+        2  4  4  0  0  0
+       [torch.FloatTensor of size 2x6]
+    """
+    out, index = gen(src, index, dim, out, dim_size, fill_value)
+    return ScatterAdd.apply(out, src, index, dim)