more doc, bugfix

torch_scatter/functions/add.py

@@ -2,7 +2,7 @@ from .utils import gen_output
 
 
 def scatter_add_(output, index, input, dim=0):
-    """Sums all values from the :attr:`input` tensor into :attr:`output` at the
+    r"""Sums all values from the :attr:`input` tensor 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
@@ -37,7 +37,7 @@ def scatter_add_(output, index, input, dim=0):
     .. testcode::
 
         from torch_scatter import scatter_add_
-        input = torch.Tensor([[2, 0, 1, 4, 3], [0, 2, 1, 3, 4]])
+        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 = torch.zeros(2, 6)
         scatter_add_(output, index, input, dim=1)
@@ -53,7 +53,7 @@ def scatter_add_(output, index, input, dim=0):
 
 
 def scatter_add(index, input, dim=0, size=None, fill_value=0):
-    """Sums all values from the :attr:`input` tensor at the indices specified
+    r"""Sums all values from the :attr:`input` tensor at the indices specified
     in the :attr:`index` tensor along an given axis :attr:`dim` (`cf.`
     :meth:`~torch_scatter.scatter_add_`).
 
@@ -85,7 +85,7 @@ def scatter_add(index, input, dim=0, size=None, fill_value=0):
     .. testcode::
 
         from torch_scatter import scatter_add
-        input = torch.Tensor([[2, 0, 1, 4, 3], [0, 2, 1, 3, 4]])
+        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_add(index, input, dim=1)
         print(output)

torch_scatter/functions/div.py

@@ -3,11 +3,89 @@ from .utils import gen_output
 
 
 def scatter_div_(output, index, input, dim=0):
-    """If multiple indices reference the same location, their
-    **contributions divide**."""
+    r"""Divides all values from the :attr:`input` tensor into :attr:`output`
+    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{output}_i = \mathrm{output}_i \cdot \prod_j
+        \frac{1}{\mathrm{input}_j}
+
+    where sum is over :math:`j` such that :math:`\mathrm{index}_j = i`.
+
+    Args:
+        output (Tensor): The destination tensor
+        index (LongTensor): The indices of elements to scatter
+        input (Tensor): The source tensor
+        dim (int, optional): The axis along which to index
+
+    :rtype: :class:`Tensor`
+
+    .. testsetup::
+
+        import torch
+
+    .. testcode::
+
+        from torch_scatter import scatter_div_
+        input =     torch.Tensor([[2, 1, 2, 4, 3], [1, 2, 2, 3, 4]])
+        index = torch.LongTensor([[4, 5, 4, 2, 3], [0, 0, 2, 2, 1]])
+        output = torch.ones(2, 6)
+        scatter_div_(output, index, input, dim=1)
+        print(output)
+
+    .. 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]
+    """
     return scatter('div', dim, output, index, input)
 
 
 def scatter_div(index, input, dim=0, size=None, fill_value=1):
+    r"""Divides all values from the :attr:`input` tensor at the indices
+    specified in the :attr:`index` tensor along an given axis :attr:`dim`
+    (`cf.` :meth:`~torch_scatter.scatter_div_` and
+    :meth:`~torch_scatter.scatter_add`).
+
+    For one-dimensional tensors, the operation computes
+
+    .. math::
+        \mathrm{output}_i = \mathrm{fill\_value} \cdot \prod_j
+        \frac{1}{\mathrm{input}_j}
+
+    where sum is over :math:`j` such that :math:`\mathrm{index}_j = i`.
+
+    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
+
+    .. testcode::
+
+        from torch_scatter import scatter_div
+        input =     torch.Tensor([[2, 1, 2, 4, 3], [1, 2, 2, 3, 4]])
+        index = torch.LongTensor([[4, 5, 4, 2, 3], [0, 0, 2, 2, 1]])
+        output = scatter_div(index, input, dim=1)
+        print(output)
+
+    .. 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]
+    """
     output = gen_output(index, input, dim, size, fill_value)
-    scatter_div_(output, index, input, dim)
+    return scatter_div_(output, index, input, dim)

torch_scatter/functions/mul.py

@@ -3,11 +3,87 @@ from .utils import gen_output
 
 
 def scatter_mul_(output, index, input, dim=0):
-    """If multiple indices reference the same location, their
-    **contributions multiply**."""
+    r"""Multiplies all values from the :attr:`input` tensor into :attr:`output`
+    at the indices specified in the :attr:`index` tensor along an given axis
+    :attr:`dim`. If multiple indices reference the same location, their
+    **contributions multiply** (`cf.` :meth:`~torch_scatter.scatter_add_`).
+
+    For one-dimensional tensors, the operation computes
+
+    .. math::
+        \mathrm{output}_i = \mathrm{output}_i \cdot \prod_j \mathrm{input}_j
+
+    where sum is over :math:`j` such that :math:`\mathrm{index}_j = i`.
+
+    Args:
+        output (Tensor): The destination tensor
+        index (LongTensor): The indices of elements to scatter
+        input (Tensor): The source tensor
+        dim (int, optional): The axis along which to index
+
+    :rtype: :class:`Tensor`
+
+    .. testsetup::
+
+        import torch
+
+    .. testcode::
+
+        from torch_scatter import scatter_mul_
+        input =     torch.Tensor([[2, 0, 3, 4, 3], [2, 3, 4, 2, 4]])
+        index = torch.LongTensor([[4, 5, 4, 2, 3], [0, 0, 2, 2, 1]])
+        output = torch.ones(2, 6)
+        scatter_mul_(output, index, input, dim=1)
+        print(output)
+
+    .. testoutput::
+
+        1  1  4  3  6  0
+        6  4  8  1  1  1
+       [torch.FloatTensor of size 2x6]
+    """
     return scatter('mul', dim, output, index, input)
 
 
 def scatter_mul(index, input, dim=0, size=None, fill_value=1):
+    r"""Multiplies all values from the :attr:`input` tensor at the indices
+    specified in the :attr:`index` tensor along an given axis :attr:`dim`
+    (`cf.` :meth:`~torch_scatter.scatter_mul_` and
+    :meth:`~torch_scatter.scatter_add`).
+
+    For one-dimensional tensors, the operation computes
+
+    .. math::
+        \mathrm{output}_i = \mathrm{fill\_value} \cdot \prod_j \mathrm{input}_j
+
+    where sum is over :math:`j` such that :math:`\mathrm{index}_j = i`.
+
+    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
+
+    .. testcode::
+
+        from torch_scatter import scatter_mul
+        input =     torch.Tensor([[2, 0, 3, 4, 3], [2, 3, 4, 2, 4]])
+        index = torch.LongTensor([[4, 5, 4, 2, 3], [0, 0, 2, 2, 1]])
+        output = scatter_mul(index, input, dim=1)
+        print(output)
+
+    .. testoutput::
+
+        1  1  4  3  6  0
+        6  4  8  1  1  1
+       [torch.FloatTensor of size 2x6]
+    """
     output = gen_output(index, input, dim, size, fill_value)
     return scatter_mul_(output, index, input, dim)

torch_scatter/functions/sub.py

@@ -2,7 +2,7 @@ from .utils import gen_output
 
 
 def scatter_sub_(output, index, input, dim=0):
-    """Subtracts all values from the :attr:`input` tensor into :attr:`output`
+    r"""Subtracts all values from the :attr:`input` tensor into :attr:`output`
     at the indices specified in the :attr:`index` tensor along an given axis
     :attr:`dim`. If multiple indices reference the same location, their
     **negated contributions add** (`cf.` :meth:`~torch_scatter.scatter_add_`).
@@ -29,7 +29,7 @@ def scatter_sub_(output, index, input, dim=0):
     .. testcode::
 
         from torch_scatter import scatter_sub_
-        input = torch.Tensor([[2, 0, 1, 4, 3], [0, 2, 1, 3, 4]])
+        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 = torch.zeros(2, 6)
         scatter_sub_(output, index, input, dim=1)
@@ -38,14 +38,14 @@ def scatter_sub_(output, index, input, dim=0):
     .. testoutput::
 
         0  0 -4 -3 -3  0
-       -2 -4 -4 -0  0  0
+       -2 -4 -4  0  0  0
        [torch.FloatTensor of size 2x6]
     """
     return output.scatter_add_(dim, index, -input)
 
 
 def scatter_sub(index, input, dim=0, size=None, fill_value=0):
-    """Subtracts all values from the :attr:`input` tensor at the indices
+    r"""Subtracts all values from the :attr:`input` tensor at the indices
     specified in the :attr:`index` tensor along an given axis :attr:`dim`
     (`cf.` :meth:`~torch_scatter.scatter_sub_` and
     :meth:`~torch_scatter.scatter_add`).
@@ -73,7 +73,7 @@ def scatter_sub(index, input, dim=0, size=None, fill_value=0):
     .. testcode::
 
         from torch_scatter import scatter_sub
-        input = torch.Tensor([[2, 0, 1, 4, 3], [0, 2, 1, 3, 4]])
+        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)