added docstring

torch_cluster/graclus.py

@@ -3,6 +3,24 @@ from .utils.ffi import graclus
 
 
 def graclus_cluster(row, col, weight=None, num_nodes=None):
+    """Greedy clustering algorithm of picking an unmarked vertex and matching
+    it with one its unmarked neighbors (that maximizes its optionally edge
+    weight).
+
+    Args:
+        row (LongTensor): Source nodes.
+        col (LongTensor): Target nodes.
+        weight (Tensor, optional): Edge weights. (default: :obj:`None`)
+        num_nodes (int, optional): The number of nodes. (default: :obj:`None`)
+
+    Examples::
+
+        >>> row = torch.LongTensor([0, 1, 1, 2])
+        >>> col = torch.LongTensor([1, 0, 2, 1])
+        >>> weight = torch.Tensor([1, 1, 1, 1])
+        >>> cluster = graclus_cluster(row, col, weight)
+    """
+
     num_nodes = row.max() + 1 if num_nodes is None else num_nodes
 
     if row.is_cuda:  # pragma: no cover

torch_cluster/grid.py

@@ -2,6 +2,25 @@ from .utils.ffi import grid
 
 
 def grid_cluster(pos, size, start=None, end=None):
+    """Voxel grid clustering algorithm, which overlays a regular grid of
+    user-defined size over the point cloud and clusters all points within a
+    voxel.
+
+    Args:
+        pos (Tensor): D-dimensional position of points.
+        size (Tensor): Size of a voxel in each dimension.
+        start (Tensor or int, optional): Start position of the grid (in each
+            dimension). (default: :obj:`None`)
+        end (Tensor or int, optional): End position of the grid (in each
+            dimension). (default: :obj:`None`)
+
+    Examples::
+
+        >>> pos = torch.Tensor([[0, 0], [11, 9], [2, 8], [2, 2], [8, 3]])
+        >>> size = torch.Tensor([5, 5])
+        >>> cluster = grid_cluster(pos, size)
+    """
+
     pos = pos.unsqueeze(-1) if pos.dim() == 1 else pos
 
     assert pos.size(1) == size.size(0), (