scatter min

Author: rusty1s

test/backward.json

@@ -5,7 +5,7 @@
 from torch.autograd import gradcheck
 import torch_scatter
 
-from .utils import devices
+from .utils import dtypes, devices, tensor
 
 funcs = ['add', 'sub', 'mul', 'div', 'mean']
 indices = [2, 0, 1, 1, 0]
@@ -20,3 +20,35 @@ def test_backward(func, device):
     op = getattr(torch_scatter, 'scatter_{}'.format(func))
     data = (src, index)
     assert gradcheck(op, data, eps=1e-6, atol=1e-4) is True
+
+
+tests = [{
+    'name': 'max',
+    'src': [1, 2, 3, 4, 5],
+    'index': [2, 0, 1, 1, 0],
+    'dim': 0,
+    'fill_value': 0,
+    'grad': [4, 8, 6],
+    'expected': [6, 0, 0, 8, 4]
+}, {
+    'name': 'min',
+    'src': [1, 2, 3, 4, 5],
+    'index': [2, 0, 1, 1, 0],
+    'dim': 0,
+    'fill_value': 3,
+    'grad': [4, 8, 6],
+    'expected': [6, 4, 8, 0, 0]
+}]
+
+
+@pytest.mark.parametrize('test,dtype,device', product(tests, dtypes, devices))
+def test_arg_backward(test, dtype, device):
+    src = tensor(test['src'], dtype, device)
+    src.requires_grad_()
+    index = tensor(test['index'], torch.long, device)
+    grad = tensor(test['grad'], dtype, device)
+
+    op = getattr(torch_scatter, 'scatter_{}'.format(test['name']))
+    out, _ = op(src, index, test['dim'], fill_value=test['fill_value'])
+    out.backward(grad)
+    assert src.grad.tolist() == test['expected']

test/forward.json

@@ -76,6 +76,38 @@
     'dim': 0,
     'fill_value': 0,
     'expected': [[3, 2.5], [3, 4]]
+}, {
+    'name': 'max',
+    'src': [[2, 0, 1, 4, 3], [0, 2, 1, 3, 4]],
+    'index': [[4, 5, 4, 2, 3], [0, 0, 2, 2, 1]],
+    'dim': -1,
+    'fill_value': 0,
+    'expected': [[0, 0, 4, 3, 2, 0], [2, 4, 3, 0, 0, 0]],
+    'expected_arg': [[-1, -1, 3, 4, 0, 1], [1, 4, 3, -1, -1, -1]]
+}, {
+    'name': 'max',
+    'src': [[5, 2], [2, 5], [4, 3], [1, 3]],
+    'index': [[0, 0], [1, 1], [1, 1], [0, 0]],
+    'dim': 0,
+    'fill_value': 0,
+    'expected': [[5, 3], [4, 5]],
+    'expected_arg': [[0, 3], [2, 1]]
+}, {
+    'name': 'min',
+    'src': [[2, 0, 1, 4, 3], [0, 2, 1, 3, 4]],
+    'index': [[4, 5, 4, 2, 3], [0, 0, 2, 2, 1]],
+    'dim': -1,
+    'fill_value': 9,
+    'expected': [[9, 9, 4, 3, 1, 0], [0, 4, 1, 9, 9, 9]],
+    'expected_arg': [[-1, -1, 3, 4, 2, 1], [0, 4, 2, -1, -1, -1]]
+}, {
+    'name': 'min',
+    'src': [[5, 2], [2, 5], [4, 3], [1, 3]],
+    'index': [[0, 0], [1, 1], [1, 1], [0, 0]],
+    'dim': 0,
+    'fill_value': 9,
+    'expected': [[1, 2], [2, 3]],
+    'expected_arg': [[3, 0], [1, 2]]
 }]
 
 
@@ -86,6 +118,10 @@ def test_forward(test, dtype, device):
     expected = tensor(test['expected'], dtype, device)
 
     op = getattr(torch_scatter, 'scatter_{}'.format(test['name']))
-    output = op(src, index, test['dim'], fill_value=test['fill_value'])
+    out = op(src, index, test['dim'], fill_value=test['fill_value'])
 
-    assert output.tolist() == expected.tolist()
+    if isinstance(out, tuple):
+        assert out[0].tolist() == expected.tolist()
+        assert out[1].tolist() == test['expected_arg']
+    else:
+        assert out.tolist() == expected.tolist()

test/test_backward.py

@@ -3,10 +3,12 @@
 from .mul import scatter_mul
 from .div import scatter_div
 from .mean import scatter_mean
+from .max import scatter_max
+from .min import scatter_min
 
 __version__ = '1.0.0'
 
 __all__ = [
     'scatter_add', 'scatter_sub', 'scatter_mul', 'scatter_div', 'scatter_mean',
-    '__version__'
+    'scatter_max', 'scatter_min', '__version__'
 ]