added more doc

torch_scatter/functions/add.py

@@ -2,8 +2,8 @@ from .utils import gen_output
 
 
 def scatter_add_(output, index, input, dim=0):
-    """Sums up all values from the tensor :attr:`input` into :attr:`output` at
-    the indices specified in the :attr:`index` tensor along an given axis
+    """Sums all values from the tensor :attr:`input` into :attr:`output` at the
+    indices specified in the :attr:`index` tensor along an given axis
     :attr:`dim`. For each value in :attr:`input`, 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
@@ -53,12 +53,12 @@ def scatter_add_(output, index, input, dim=0):
 
 
 def scatter_add(index, input, dim=0, size=None, fill_value=0):
-    """Sums ap all values from the tensor :attr:`input` at the indices
-    specified in the :attr:`index` tensor along an given axis :attr:`dim`.
-    The output size at dimension :attr:`dim` is given by :attr:`size` and must
-    be at least size `index.max(dim) - 1`. If :attr:`size` is not given, a
-    minimal sized output tensor is returned. The output tensor is prefilled
-    with the specified value from :attr:`fill_value`.
+    """Sums all values from the tensor :attr:`input` at the indices specified
+    in the :attr:`index` tensor along an given axis :attr:`dim`. The output
+    size at dimension :attr:`dim` is given by :attr:`size` and must be at least
+    size `index.max(dim) - 1`. If :attr:`size` is not given, a minimal sized
+    output tensor is returned. The output tensor is prefilled with the
+    specified value from :attr:`fill_value`.
 
     For one-dimensional tensors, the operation computes
 

torch_scatter/functions/div.py

@@ -4,7 +4,7 @@ from .utils import gen_output
 
 def scatter_div_(output, index, input, dim=0):
     """If multiple indices reference the same location, their
-    contributions divide."""
+    **contributions divide**."""
     return scatter('div', dim, output, index, input)
 
 

torch_scatter/functions/max.py

@@ -3,8 +3,8 @@ from .utils import gen_filled_tensor, gen_output
 
 
 def scatter_max_(output, index, input, dim=0):
-    """If multiple indices reference the same location, the maximal
-    contribution gets taken.
+    """If multiple indices reference the same location, their **contribution
+    maximize**.
 
     :rtype: (:class:`Tensor`, :class:`LongTensor`)
     """

torch_scatter/functions/mean.py

@@ -4,7 +4,7 @@ from .utils import gen_filled_tensor, gen_output
 
 def scatter_mean_(output, index, input, dim=0):
     """If multiple indices reference the same location, their
-    contributions average."""
+    **contributions average**."""
     num_output = gen_filled_tensor(output, output.size(), fill_value=0)
     scatter('mean', dim, output, index, input, num_output)
     num_output[num_output == 0] = 1

torch_scatter/functions/min.py

@@ -3,8 +3,8 @@ from .utils import gen_filled_tensor, gen_output
 
 
 def scatter_min_(output, index, input, dim=0):
-    """If multiple indices reference the same location, the minimal
-    contribution gets taken."""
+    """If multiple indices reference the same location, their **contribution
+    minimize**."""
     arg_output = gen_filled_tensor(index, output.size(), fill_value=-1)
     return scatter('min', dim, output, index, input, arg_output)
 

torch_scatter/functions/mul.py

@@ -4,7 +4,7 @@ from .utils import gen_output
 
 def scatter_mul_(output, index, input, dim=0):
     """If multiple indices reference the same location, their
-    contributions multiply."""
+    **contributions multiply**."""
     return scatter('mul', dim, output, index, input)
 
 

torch_scatter/functions/sub.py

@@ -2,11 +2,55 @@ from .utils import gen_output
 
 
 def scatter_sub_(output, index, input, dim=0):
-    """If multiple indices reference the same location, their negated
-    contributions add."""
+    """If multiple indices reference the same location, their **negated
+    contributions add**."""
     return output.scatter_add_(dim, index, -input)
 
 
 def scatter_sub(index, input, dim=0, size=None, fill_value=0):
+    """Subtracts all values from the tensor :attr:`input` at the indices
+    specified in the :attr:`index` tensor along an given axis :attr:`dim`. The
+    output size at dimension :attr:`dim` is given by :attr:`size` and must be
+    at least size `index.max(dim) - 1`. If :attr:`size` is not given, a minimal
+    sized output tensor is returned. The output tensor is prefilled with the
+    specified value from :attr:`fill_value`.
+
+    For one-dimensional tensors, the operation computes
+
+    .. math::
+        \mathrm{output}_i = \mathrm{fill\_value} - \sum_j \mathrm{input}_j
+
+    where sum is over :math:`j` such that :math:`\mathrm{index}_j = i`.
+
+    A more detailed explanation is described in
+    :meth:`~torch_scatter.scatter_sub_`.
+
+    Args:
+        index (LongTensor): The indices of elements to scatter
+        input (Tensor): The source tensor
+        dim (int, optional): The axis along which to index
+        size (int, optional): Output size at dimension :attr:`dim`
+        fill_value (int, optional): Initial filling of output tensor
+
+    :rtype: :class:`Tensor`
+
+    .. testsetup::
+
+        import torch
+        from torch_scatter import scatter_sub
+
+    .. testcode::
+
+        input = torch.Tensor([[2, 0, 1, 4, 3], [0, 2, 1, 3, 4]])
+        index = torch.LongTensor([[4, 5, 4, 2, 3], [0, 0, 2, 2, 1]])
+        output = scatter_sub(index, input, dim=1)
+        print(output)
+
+    .. testoutput::
+
+        0  0 -4 -3 -3  0
+       -2 -4 -4  0  0  0
+       [torch.FloatTensor of size 2x6]
+    """
     output = gen_output(index, input, dim, size, fill_value)
     return scatter_sub_(output, index, input, dim)