scatter add doc

docs/source/functions/add.rst

+Scatter Add
+===========
+
+.. automodule:: torch_scatter
+
+.. autofunction:: scatter_add_
+.. autofunction:: scatter_add

docs/source/index.rst

@@ -6,7 +6,9 @@ PyTorch Scatter documentation
 .. toctree::
    :glob:
    :maxdepth: 1
-   :caption: PackageReference
+   :caption: Package reference
+
+   functions/*
 
 Indices and tables
 ==================

torch_scatter/functions/add.py

@@ -7,8 +7,8 @@ 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
     :attr:`dim`. For each value in :attr:`input`, its output index is specified
-    by its index in :attr:`input` for dimension != :attr:`dim` and by the
-    corresponding value in :attr:`index` for dimension = :attr:`dim`. If
+    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
     multiple indices reference the same location, their contributions add.
 
     If :attr:`input` and :attr:`index` are n-dimensional tensors with size
@@ -18,8 +18,11 @@ def scatter_add_(output, index, input, dim=0):
     values of :attr:`index` must be between `0` and `output.size(dim) - 1`.
 
     For one-dimensional tensors, the operation computes
-    :math:`output_i = output_i + \sum_j input_j`, where sum is over
-    :math:`j` such that :math:`index_j = i`.
+
+    .. math::
+        \mathrm{output}_i = \mathrm{output}_i + \sum_j \mathrm{input}_j
+
+    where sum is over :math:`j` such that :math:`\mathrm{index}_j = i`.
 
     Args:
         output (Tensor): The destination tensor
@@ -28,10 +31,12 @@ def scatter_add_(output, index, input, dim=0):
         dim (int, optional): The axis along which to index
 
     Example::
+
         >> input = torch.Tensor([[2, 0, 1, 4, 3], [0,2, 1, 3, 4]])
         >> index = torch.LongTensor([[4, 5, 2, 3], [0, 0, 2, 2, 1]])
         >> output = torch.zeros(2, 6)
         >> scatter_add_(output, index, input, dim=1)
+
         0  0  4  3  3  0
         2  4  4  0  0  0
         [torch.FloatTensor of size 2x6]
@@ -46,14 +51,18 @@ def scatter_add(index, input, dim=0, size=None, fill_value=0):
     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 initially
-    filled with the specified value at :attr:`fill_value`.
+    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:`output_i = fill_value + \sum_j input_j`, where sum is over
-    :math:`j` such that :math:`index_j = i`.
 
-    A more detailed explanation is described in :meth:`~scatter_add_`.
+    .. 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_add_`.
 
     Args:
         index (LongTensor): The indices of elements to scatter