doc fixes

build.py

 import os.path as osp
-import shutil
 import subprocess
 
 import torch
 from torch.utils.ffi import create_extension
 
-if osp.exists('build'):
-    shutil.rmtree('build')
-
 headers = ['torch_scatter/src/cpu.h']
 sources = ['torch_scatter/src/cpu.c']
 include_dirs = ['torch_scatter/src']

docs/source/conf.py

-import os
-import sys
 import datetime
 import sphinx_rtd_theme
 import doctest
 
-sys.path.insert(0, os.path.abspath('../..'))
-
 from torch_scatter import __version__  # noqa
 
 extensions = [

docs/source/functions/add.rst

@@ -3,5 +3,4 @@ Scatter Add
 
 .. automodule:: torch_scatter
 
-.. autofunction:: scatter_add_
 .. autofunction:: scatter_add

docs/source/functions/div.rst

@@ -3,5 +3,4 @@ Scatter Div
 
 .. automodule:: torch_scatter
 
-.. autofunction:: scatter_div_
 .. autofunction:: scatter_div

docs/source/functions/max.rst

@@ -3,5 +3,4 @@ Scatter Max
 
 .. automodule:: torch_scatter
 
-.. autofunction:: scatter_max_
 .. autofunction:: scatter_max

docs/source/functions/mean.rst

@@ -3,5 +3,4 @@ Scatter Mean
 
 .. automodule:: torch_scatter
 
-.. autofunction:: scatter_mean_
 .. autofunction:: scatter_mean

docs/source/functions/min.rst

@@ -3,5 +3,4 @@ Scatter Min
 
 .. automodule:: torch_scatter
 
-.. autofunction:: scatter_min_
 .. autofunction:: scatter_min

docs/source/functions/mul.rst

@@ -3,5 +3,4 @@ Scatter Mul
 
 .. automodule:: torch_scatter
 
-.. autofunction:: scatter_mul_
 .. autofunction:: scatter_mul

docs/source/functions/sub.rst

@@ -3,5 +3,4 @@ Scatter Sub
 
 .. automodule:: torch_scatter
 
-.. autofunction:: scatter_sub_
 .. autofunction:: scatter_sub

torch_scatter/add.py

@@ -50,7 +50,7 @@ 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 :math:`\sum` is over :math:`j` such that
+    where :math:`\sum_j` is over :math:`j` such that
     :math:`\mathrm{index}_j = i`.
 
     Args:
@@ -75,17 +75,19 @@ def scatter_add(src, index, dim=-1, out=None, dim_size=None, fill_value=0):
     .. 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]
+       tensor([[ 0,  0,  4,  3,  3,  0],
+               [ 2,  4,  4,  0,  0,  0]])
     """
     src, out, index, dim = gen(src, index, dim, out, dim_size, fill_value)
     return ScatterAdd.apply(out, src, index, dim)

torch_scatter/div.py

@@ -48,7 +48,7 @@ def scatter_div(src, index, dim=-1, out=None, dim_size=None, fill_value=1):
         \mathrm{out}_i = \mathrm{out}_i \cdot \prod_j
         \frac{1}{\mathrm{src}_j}
 
-    where :math:`\prod` is over :math:`j` such that
+    where :math:`\prod_j` is over :math:`j` such that
     :math:`\mathrm{index}_j = i`.
 
     Args:
@@ -73,17 +73,19 @@ def scatter_div(src, index, dim=-1, out=None, dim_size=None, fill_value=1):
     .. testcode::
 
         from torch_scatter import scatter_div
-        src = torch.tensor([[2, 0, 3, 4, 3], [2, 3, 4, 2, 4]])
+
+        src = torch.tensor([[2, 1, 1, 4, 2], [1, 2, 1, 2, 4]]).float()
         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]
+       tensor([[ 1.0000,  1.0000,  0.2500,  0.5000,  0.5000,  1.0000],
+               [ 0.5000,  0.2500,  0.5000,  1.0000,  1.0000,  1.0000]])
     """
     src, out, index, dim = gen(src, index, dim, out, dim_size, fill_value)
     return ScatterDiv.apply(out, src, index, dim)

torch_scatter/max.py

@@ -53,7 +53,7 @@ def scatter_max(src, index, dim=-1, out=None, dim_size=None, fill_value=0):
     .. math::
         \mathrm{out}_i = \max(\mathrm{out}_i, \max_j(\mathrm{src}_j))
 
-    where :math:`\max` is over :math:`j` such that
+    where :math:`\max_j` is over :math:`j` such that
     :math:`\mathrm{index}_j = i`.
 
     Args:
@@ -78,23 +78,22 @@ def scatter_max(src, index, dim=-1, out=None, dim_size=None, fill_value=0):
     .. testcode::
 
         from torch_scatter import scatter_max
+
         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_max(src, index, out=out)
+
+        out, argmax = scatter_max(src, index, out=out)
+
         print(out)
+        print(argmax)
 
     .. testoutput::
 
-       (
-        0  0  4  3  2  0
-        2  4  3  0  0  0
-       [torch.FloatTensor of size 2x6]
-       ,
-       -1 -1  3  4  0  1
-        1  4  3 -1 -1 -1
-       [torch.LongTensor of size 2x6]
-       )
+       tensor([[ 0,  0,  4,  3,  2,  0],
+               [ 2,  4,  3,  0,  0,  0]])
+       tensor([[-1, -1,  3,  4,  0,  1],
+               [ 1,  4,  3, -1, -1, -1]])
     """
     src, out, index, dim = gen(src, index, dim, out, dim_size, fill_value)
     return ScatterMax.apply(out, src, index, dim)

torch_scatter/mean.py

@@ -51,8 +51,9 @@ 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 :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`.
+    where :math:`\sum_j` is over :math:`j` such that
+    :math:`\mathrm{index}_j = i`. :math:`N_i` indicates the number of indices
+    referencing :math:`i`.
 
     Args:
         src (Tensor): The source tensor.
@@ -76,17 +77,19 @@ def scatter_mean(src, index, dim=-1, out=None, dim_size=None, fill_value=0):
     .. testcode::
 
         from torch_scatter import scatter_mean
-        src = torch.tensor([[2, 0, 1, 4, 3], [0, 2, 1, 3, 4]])
+
+        src = torch.tensor([[2, 0, 1, 4, 3], [0, 2, 1, 3, 4]]).float()
         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)
+
         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
-       [torch.FloatTensor of size 2x6]
+       tensor([[ 0.0000,  0.0000,  4.0000,  3.0000,  1.5000,  0.0000],
+               [ 1.0000,  4.0000,  2.0000,  0.0000,  0.0000,  0.0000]])
     """
     src, out, index, dim = gen(src, index, dim, out, dim_size, fill_value)
     return ScatterMean.apply(out, src, index, dim)

torch_scatter/min.py

@@ -53,7 +53,7 @@ def scatter_min(src, index, dim=-1, out=None, dim_size=None, fill_value=0):
     .. math::
         \mathrm{out}_i = \min(\mathrm{out}_i, \min_j(\mathrm{src}_j))
 
-    where :math:`\min` is over :math:`j` such that
+    where :math:`\min_j` is over :math:`j` such that
     :math:`\mathrm{index}_j = i`.
 
     Args:
@@ -77,24 +77,23 @@ def scatter_min(src, index, dim=-1, out=None, dim_size=None, fill_value=0):
 
     .. testcode::
 
-        from torch_scatter import scatter_mean
+        from torch_scatter import scatter_min
+
         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_min(src, index, out=out)
+
+        out, argmax = scatter_min(src, index, out=out)
+
         print(out)
+        print(argmax)
 
     .. testoutput::
 
-       (
-        0  0 -4 -3 -2  0
-       -2 -4 -3  0  0  0
-       [torch.FloatTensor of size 2x6]
-       ,
-       -1 -1  3  4  0  1
-        1  4  3 -1 -1 -1
-       [torch.LongTensor of size 2x6]
-       )
+       tensor([[ 0,  0, -4, -3, -2,  0],
+               [-2, -4, -3,  0,  0,  0]])
+       tensor([[-1, -1,  3,  4,  0,  1],
+               [ 1,  4,  3, -1, -1, -1]])
     """
     src, out, index, dim = gen(src, index, dim, out, dim_size, fill_value)
     return ScatterMin.apply(out, src, index, dim)

torch_scatter/mul.py

@@ -47,7 +47,7 @@ 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 :math:`\prod` is over :math:`j` such that
+    where :math:`\prod_j` is over :math:`j` such that
     :math:`\mathrm{index}_j = i`.
 
     Args:
@@ -72,17 +72,19 @@ def scatter_mul(src, index, dim=-1, out=None, dim_size=None, fill_value=1):
     .. testcode::
 
         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_ones((2, 6))
+
         out = scatter_mul(src, index, out=out)
+
         print(out)
 
     .. testoutput::
 
-        1  1  4  3  6  0
-        6  4  8  1  1  1
-       [torch.FloatTensor of size 2x6]
+       tensor([[ 1,  1,  4,  3,  6,  0],
+               [ 6,  4,  8,  1,  1,  1]])
     """
     src, out, index, dim = gen(src, index, dim, out, dim_size, fill_value)
     return ScatterMul.apply(out, src, index, dim)

torch_scatter/sub.py

@@ -22,7 +22,7 @@ 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 :math:`\sum` is over :math:`j` such that
+    where :math:`\sum_j` is over :math:`j` such that
     :math:`\mathrm{index}_j = i`.
 
     Args:
@@ -47,16 +47,18 @@ def scatter_sub(src, index, dim=-1, out=None, dim_size=None, fill_value=0):
     .. testcode::
 
         from torch_scatter import scatter_sub
+
         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_sub(src, index, out=out)
+
         print(out)
 
     .. testoutput::
 
-        0  0 -4 -3 -3  0
-       -2 -4 -4  0  0  0
-       [torch.FloatTensor of size 2x6]
+       tensor([[ 0,  0, -4, -3, -3,  0],
+               [-2, -4, -4,  0,  0,  0]])
     """
     return scatter_add(src.neg(), index, dim, out, dim_size, fill_value)