scatter div

Author: rusty1s

test/test_backward.py

@@ -7,7 +7,7 @@
 
 from .utils import devices
 
-funcs = ['add', 'sub', 'mul', 'mean']
+funcs = ['add', 'sub', 'mul', 'div', 'mean']
 indices = [2, 0, 1, 1, 0]
 
 

test/test_forward.py

@@ -48,6 +48,20 @@
     'dim': 0,
     'fill_value': 1,
     'expected': [[5, 6], [8, 15]]
+}, {
+    'name': 'div',
+    'src': [[2, 1, 1, 4, 2], [1, 2, 1, 2, 4]],
+    'index': [[4, 5, 4, 2, 3], [0, 0, 2, 2, 1]],
+    'dim': -1,
+    'fill_value': 1,
+    'expected': [[1, 1, 0.25, 0.5, 0.5, 1], [0.5, 0.25, 0.5, 1, 1, 1]]
+}, {
+    'name': 'div',
+    'src': [[4, 2], [2, 1], [4, 2], [1, 2]],
+    'index': [[0, 0], [1, 1], [1, 1], [0, 0]],
+    'dim': 0,
+    'fill_value': 1,
+    'expected': [[0.25, 0.25], [0.125, 0.5]]
 }, {
     'name': 'mean',
     'src': [[2, 0, 1, 4, 3], [0, 2, 1, 3, 4]],

torch_scatter/__init__.py

@@ -1,10 +1,12 @@
 from .add import scatter_add
 from .sub import scatter_sub
 from .mul import scatter_mul
+from .div import scatter_div
 from .mean import scatter_mean
 
 __version__ = '1.0.0'
 
 __all__ = [
-    'scatter_add', 'scatter_sub', 'scatter_mul', 'scatter_mean', '__version__'
+    'scatter_add', 'scatter_sub', 'scatter_mul', 'scatter_div', 'scatter_mean',
+    '__version__'
 ]

torch_scatter/add.py

@@ -50,7 +50,8 @@ def scatter_add(src, index, dim=-1, out=None, dim_size=None, fill_value=0):
     .. 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`.
+    where :math:`\sum` is over :math:`j` such that
+    :math:`\mathrm{index}_j = i`.
 
     Args:
         src (Tensor): The source tensor.

torch_scatter/div.py

@@ -0,0 +1,89 @@
+from torch.autograd import Function
+
+from .utils.ffi import get_func
+from .utils.gen import gen
+
+
+class ScatterDiv(Function):
+    @staticmethod
+    def forward(ctx, out, src, index, dim):
+        func = get_func('scatter_div', src)
+        func(dim, out, index, src)
+
+        ctx.mark_dirty(out)
+        ctx.save_for_backward(out, src, index)
+
+        return out
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        out, src, index = ctx.saved_variables
+
+        grad_src = None
+        if ctx.needs_input_grad[1]:
+            grad_src = -(out * grad_out)[index] / src
+
+        return None, grad_src, None, None
+
+
+def scatter_div(src, index, dim=-1, out=None, dim_size=None, fill_value=1):
+    r"""
+    |
+
+    .. image:: https://raw.githubusercontent.com/rusty1s/pytorch_scatter/
+            master/docs/source/_figures/div.svg?sanitize=true
+        :align: center
+        :width: 400px
+
+    |
+
+    Divides 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`.If multiple indices reference the same location, their
+    **contributions divide** (`cf.` :meth:`~torch_scatter.scatter_add`).
+
+    For one-dimensional tensors, the operation computes
+
+    .. math::
+        \mathrm{out}_i = \mathrm{out}_i \cdot \prod_j
+        \frac{1}{\mathrm{src}_j}
+
+    where :math:`\prod` 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_div
+        src = torch.tensor([[2, 0, 3, 4, 3], [2, 3, 4, 2, 4]])
+        index = torch.tensor([[4, 5, 4, 2, 3], [0, 0, 2, 2, 1]])
+        out = src.new_ones((2, 6))
+        out = scatter_div(src, index, out=out)
+        print(out)
+
+    .. testoutput::
+
+        1.0000  1.0000  0.2500  0.3333  0.2500  1.0000
+        0.5000  0.2500  0.1667  1.0000  1.0000  1.0000
+       [torch.FloatTensor of size 2x6]
+    """
+    src, out, index, dim = gen(src, index, dim, out, dim_size, fill_value)
+    return ScatterDiv.apply(out, src, index, dim)

torch_scatter/mean.py

@@ -51,8 +51,8 @@ def scatter_mean(src, index, dim=-1, out=None, dim_size=None, fill_value=0):
         \mathrm{out}_i = \mathrm{out}_i + \frac{1}{N_i} \cdot
         \sum_j \mathrm{src}_j
 
-    where sum is over :math:`j` such that :math:`\mathrm{index}_j = i` and
-    :math:`N_i` indicates the number of indices referencing :math:`i`.
+    where :math:`\sum` is over :math:`j` such that :math:`\mathrm{index}_j = i`
+    add :math:`N_i` indicates the number of indices referencing :math:`i`.
 
     Args:
         src (Tensor): The source tensor.

torch_scatter/mul.py

@@ -10,7 +10,6 @@ def forward(ctx, out, src, index, dim):
         func = get_func('scatter_mul', src)
         func(dim, out, index, src)
 
-        ctx.dim = dim
         ctx.mark_dirty(out)
         ctx.save_for_backward(out, src, index)
 
@@ -48,7 +47,8 @@ def scatter_mul(src, index, dim=-1, out=None, dim_size=None, fill_value=1):
     .. math::
         \mathrm{out}_i = \mathrm{out}_i \cdot \prod_j \mathrm{src}_j
 
-    where sum is over :math:`j` such that :math:`\mathrm{index}_j = i`.
+    where :math:`\prod` is over :math:`j` such that
+    :math:`\mathrm{index}_j = i`.
 
     Args:
         src (Tensor): The source tensor.
@@ -71,17 +71,17 @@ def scatter_mul(src, index, dim=-1, out=None, dim_size=None, fill_value=1):
 
     .. testcode::
 
-        from torch_scatter import scatter_mean
-        src = torch.tensor([[2, 0, 1, 4, 3], [0, 2, 1, 3, 4]])
+        from torch_scatter import scatter_mul
+        src = torch.tensor([[2, 0, 3, 4, 3], [2, 3, 4, 2, 4]])
         index = torch.tensor([[4, 5, 4, 2, 3], [0, 0, 2, 2, 1]])
-        out = src.new_zeros((2, 6))
-        out = scatter_mean(src, index, out=out)
+        out = src.new_ones((2, 6))
+        out = scatter_mul(src, index, out=out)
         print(out)
 
     .. testoutput::
 
-        0.0000  0.0000  4.0000  3.0000  1.5000  0.0000
-        1.0000  4.0000  2.0000  0.0000  0.0000  0.0000
+        1  1  4  3  6  0
+        6  4  8  1  1  1
        [torch.FloatTensor of size 2x6]
     """
     src, out, index, dim = gen(src, index, dim, out, dim_size, fill_value)

torch_scatter/sub.py

@@ -22,7 +22,8 @@ def scatter_sub(src, index, dim=-1, out=None, dim_size=None, fill_value=0):
     .. 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`.
+    where :math:`\sum` is over :math:`j` such that
+    :math:`\mathrm{index}_j = i`.
 
     Args:
         src (Tensor): The source tensor.