added docstring

torch_scatter/functions/__init__.py

@@ -5,8 +5,34 @@ from .utils import gen_filled_tensor, gen_output
 
 
 def scatter_add_(output, index, input, dim=0):
-    """If multiple indices reference the same location, their contributions
-    add."""
+    """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
+    multiple indices reference the same location, their contributions add.
+
+    If :attr:`input` and :attr:`index` are n-dimensional tensors with equal
+    size :math:`(x_0, ..., x_{i-1}, x_i, x_{i+1}, ..., x_{n-1})` and
+    :attr:`dim` = i, then :attr:`output` must be an n-dimensional tensor with
+    size :math:`(x_0, ..., x_{i-1}, y, x_{i+1}, ..., x_{n-1})`. Moreover, the
+    values of :attr:`index` must be between `0` and `output.size(dim) - 1`.
+
+    Args:
+        output (Tensor): The destination tensor
+        index (LongTensor): The indices of elements to scatter
+        input (Tensor): The source tensor
+        dim (int): 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]
+    """
     return output.scatter_add_(dim, index, input)