Initial commit

torch_cluster.egg-info/top_level.txt

+torch_cluster

torch_cluster/__init__.py

+import importlib
+import os.path as osp
+
+import torch
+
+__version__ = '1.6.0'
+
+for library in [
+        '_version', '_grid', '_graclus', '_fps', '_rw', '_sampler', '_nearest',
+        '_knn', '_radius'
+]:
+    hip_spec = importlib.machinery.PathFinder().find_spec(
+        f'{library}_hip', [osp.dirname(__file__)])
+    cpu_spec = importlib.machinery.PathFinder().find_spec(
+        f'{library}_cpu', [osp.dirname(__file__)])
+    spec = hip_spec or cpu_spec
+    if spec is not None:
+        torch.ops.load_library(spec.origin)
+    else:  # pragma: no cover
+        raise ImportError(f"Could not find module '{library}_cpu' in "
+                          f"{osp.dirname(__file__)}")
+
+cuda_version = torch.ops.torch_cluster.cuda_version()
+if torch.cuda.is_available() and cuda_version != -1:  # pragma: no cover
+    if cuda_version < 10000:
+        major, minor = int(str(cuda_version)[0]), int(str(cuda_version)[2])
+    else:
+        major, minor = int(str(cuda_version)[0:2]), int(str(cuda_version)[3])
+
+from .fps import fps  # noqa
+from .graclus import graclus_cluster  # noqa
+from .grid import grid_cluster  # noqa
+from .knn import knn, knn_graph  # noqa
+from .nearest import nearest  # noqa
+from .radius import radius, radius_graph  # noqa
+from .rw import random_walk  # noqa
+from .sampler import neighbor_sampler  # noqa
+
+__all__ = [
+    'graclus_cluster',
+    'grid_cluster',
+    'fps',
+    'nearest',
+    'knn',
+    'knn_graph',
+    'radius',
+    'radius_graph',
+    'random_walk',
+    'neighbor_sampler',
+    '__version__',
+]

torch_cluster/fps.py

+from typing import Optional
+
+import torch
+from torch import Tensor
+
+
+@torch.jit._overload  # noqa
+def fps(src, batch=None, ratio=None, random_start=True):  # noqa
+    # type: (Tensor, Optional[Tensor], Optional[float], bool) -> Tensor
+    pass  # pragma: no cover
+
+
+@torch.jit._overload  # noqa
+def fps(src, batch=None, ratio=None, random_start=True):  # noqa
+    # type: (Tensor, Optional[Tensor], Optional[Tensor], bool) -> Tensor
+    pass  # pragma: no cover
+
+
+def fps(src: torch.Tensor, batch=None, ratio=None, random_start=True):  # noqa
+    r""""A sampling algorithm from the `"PointNet++: Deep Hierarchical Feature
+    Learning on Point Sets in a Metric Space"
+    <https://arxiv.org/abs/1706.02413>`_ paper, which iteratively samples the
+    most distant point with regard to the rest points.
+
+    Args:
+        src (Tensor): Point feature matrix
+            :math:`\mathbf{X} \in \mathbb{R}^{N \times F}`.
+        batch (LongTensor, optional): Batch vector
+            :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each
+            node to a specific example. (default: :obj:`None`)
+        ratio (float or Tensor, optional): Sampling ratio.
+            (default: :obj:`0.5`)
+        random_start (bool, optional): If set to :obj:`False`, use the first
+            node in :math:`\mathbf{X}` as starting node. (default: obj:`True`)
+
+    :rtype: :class:`LongTensor`
+
+
+    .. code-block:: python
+
+        import torch
+        from torch_cluster import fps
+
+        src = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]])
+        batch = torch.tensor([0, 0, 0, 0])
+        index = fps(src, batch, ratio=0.5)
+    """
+
+    r: Optional[Tensor] = None
+    if ratio is None:
+        r = torch.tensor(0.5, dtype=src.dtype, device=src.device)
+    elif isinstance(ratio, float):
+        r = torch.tensor(ratio, dtype=src.dtype, device=src.device)
+    else:
+        r = ratio
+    assert r is not None
+
+    if batch is not None:
+        assert src.size(0) == batch.numel()
+        batch_size = int(batch.max()) + 1
+
+        deg = src.new_zeros(batch_size, dtype=torch.long)
+        deg.scatter_add_(0, batch, torch.ones_like(batch))
+
+        ptr = deg.new_zeros(batch_size + 1)
+        torch.cumsum(deg, 0, out=ptr[1:])
+    else:
+        ptr = torch.tensor([0, src.size(0)], device=src.device)
+
+    return torch.ops.torch_cluster.fps(src, ptr, r, random_start)

torch_cluster/graclus.py

+from typing import Optional
+
+import torch
+
+
+@torch.jit.script
+def graclus_cluster(row: torch.Tensor, col: torch.Tensor,
+                    weight: Optional[torch.Tensor] = None,
+                    num_nodes: Optional[int] = None) -> torch.Tensor:
+    """A greedy clustering algorithm of picking an unmarked vertex and matching
+    it with one its unmarked neighbors (that maximizes its 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`)
+
+    :rtype: :class:`LongTensor`
+
+    .. code-block:: python
+
+        import torch
+        from torch_cluster import graclus_cluster
+
+        row = torch.tensor([0, 1, 1, 2])
+        col = torch.tensor([1, 0, 2, 1])
+        weight = torch.Tensor([1, 1, 1, 1])
+        cluster = graclus_cluster(row, col, weight)
+    """
+
+    if num_nodes is None:
+        num_nodes = max(int(row.max()), int(col.max())) + 1
+
+    # Remove self-loops.
+    mask = row != col
+    row, col = row[mask], col[mask]
+
+    if weight is not None:
+        weight = weight[mask]
+
+    # Randomly shuffle nodes.
+    if weight is None:
+        perm = torch.randperm(row.size(0), dtype=torch.long, device=row.device)
+        row, col = row[perm], col[perm]
+
+    # To CSR.
+    perm = torch.argsort(row)
+    row, col = row[perm], col[perm]
+
+    if weight is not None:
+        weight = weight[perm]
+
+    deg = row.new_zeros(num_nodes)
+    deg.scatter_add_(0, row, torch.ones_like(row))
+    rowptr = row.new_zeros(num_nodes + 1)
+    torch.cumsum(deg, 0, out=rowptr[1:])
+
+    return torch.ops.torch_cluster.graclus(rowptr, col, weight)

torch_cluster/grid.py

+from typing import Optional
+
+import torch
+
+
+@torch.jit.script
+def grid_cluster(pos: torch.Tensor, size: torch.Tensor,
+                 start: Optional[torch.Tensor] = None,
+                 end: Optional[torch.Tensor] = None) -> torch.Tensor:
+    """A clustering algorithm, which overlays a regular grid of user-defined
+    size over a 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, optional): Start position of the grid (in each
+            dimension). (default: :obj:`None`)
+        end (Tensor, optional): End position of the grid (in each
+            dimension). (default: :obj:`None`)
+
+    :rtype: :class:`LongTensor`
+
+    .. code-block:: python
+
+        import torch
+        from torch_cluster import grid_cluster
+
+        pos = torch.Tensor([[0, 0], [11, 9], [2, 8], [2, 2], [8, 3]])
+        size = torch.Tensor([5, 5])
+        cluster = grid_cluster(pos, size)
+    """
+    return torch.ops.torch_cluster.grid(pos, size, start, end)

torch_cluster/knn.py

+from typing import Optional
+
+import torch
+
+
+@torch.jit.script
+def knn(x: torch.Tensor, y: torch.Tensor, k: int,
+        batch_x: Optional[torch.Tensor] = None,
+        batch_y: Optional[torch.Tensor] = None, cosine: bool = False,
+        num_workers: int = 1) -> torch.Tensor:
+    r"""Finds for each element in :obj:`y` the :obj:`k` nearest points in
+    :obj:`x`.
+
+    Args:
+        x (Tensor): Node feature matrix
+            :math:`\mathbf{X} \in \mathbb{R}^{N \times F}`.
+        y (Tensor): Node feature matrix
+            :math:`\mathbf{X} \in \mathbb{R}^{M \times F}`.
+        k (int): The number of neighbors.
+        batch_x (LongTensor, optional): Batch vector
+            :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each
+            node to a specific example. :obj:`batch_x` needs to be sorted.
+            (default: :obj:`None`)
+        batch_y (LongTensor, optional): Batch vector
+            :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^M`, which assigns each
+            node to a specific example. :obj:`batch_y` needs to be sorted.
+            (default: :obj:`None`)
+        cosine (boolean, optional): If :obj:`True`, will use the Cosine
+            distance instead of the Euclidean distance to find nearest
+            neighbors. (default: :obj:`False`)
+        num_workers (int): Number of workers to use for computation. Has no
+            effect in case :obj:`batch_x` or :obj:`batch_y` is not
+            :obj:`None`, or the input lies on the GPU. (default: :obj:`1`)
+
+    :rtype: :class:`LongTensor`
+
+    .. code-block:: python
+
+        import torch
+        from torch_cluster import knn
+
+        x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]])
+        batch_x = torch.tensor([0, 0, 0, 0])
+        y = torch.Tensor([[-1, 0], [1, 0]])
+        batch_y = torch.tensor([0, 0])
+        assign_index = knn(x, y, 2, batch_x, batch_y)
+    """
+
+    x = x.view(-1, 1) if x.dim() == 1 else x
+    y = y.view(-1, 1) if y.dim() == 1 else y
+    x, y = x.contiguous(), y.contiguous()
+
+    batch_size = 1
+    if batch_x is not None:
+        assert x.size(0) == batch_x.numel()
+        batch_size = int(batch_x.max()) + 1
+    if batch_y is not None:
+        assert y.size(0) == batch_y.numel()
+        batch_size = max(batch_size, int(batch_y.max()) + 1)
+
+    ptr_x: Optional[torch.Tensor] = None
+    ptr_y: Optional[torch.Tensor] = None
+    if batch_size > 1:
+        assert batch_x is not None
+        assert batch_y is not None
+        arange = torch.arange(batch_size + 1, device=x.device)
+        ptr_x = torch.bucketize(arange, batch_x)
+        ptr_y = torch.bucketize(arange, batch_y)
+
+    return torch.ops.torch_cluster.knn(x, y, ptr_x, ptr_y, k, cosine,
+                                       num_workers)
+
+
+@torch.jit.script
+def knn_graph(x: torch.Tensor, k: int, batch: Optional[torch.Tensor] = None,
+              loop: bool = False, flow: str = 'source_to_target',
+              cosine: bool = False, num_workers: int = 1) -> torch.Tensor:
+    r"""Computes graph edges to the nearest :obj:`k` points.
+
+    Args:
+        x (Tensor): Node feature matrix
+            :math:`\mathbf{X} \in \mathbb{R}^{N \times F}`.
+        k (int): The number of neighbors.
+        batch (LongTensor, optional): Batch vector
+            :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each
+            node to a specific example. :obj:`batch` needs to be sorted.
+            (default: :obj:`None`)
+        loop (bool, optional): If :obj:`True`, the graph will contain
+            self-loops. (default: :obj:`False`)
+        flow (string, optional): The flow direction when used in combination
+            with message passing (:obj:`"source_to_target"` or
+            :obj:`"target_to_source"`). (default: :obj:`"source_to_target"`)
+        cosine (boolean, optional): If :obj:`True`, will use the Cosine
+            distance instead of Euclidean distance to find nearest neighbors.
+            (default: :obj:`False`)
+        num_workers (int): Number of workers to use for computation. Has no
+            effect in case :obj:`batch` is not :obj:`None`, or the input lies
+            on the GPU. (default: :obj:`1`)
+
+    :rtype: :class:`LongTensor`
+
+    .. code-block:: python
+
+        import torch
+        from torch_cluster import knn_graph
+
+        x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]])
+        batch = torch.tensor([0, 0, 0, 0])
+        edge_index = knn_graph(x, k=2, batch=batch, loop=False)
+    """
+
+    assert flow in ['source_to_target', 'target_to_source']
+    edge_index = knn(x, x, k if loop else k + 1, batch, batch, cosine,
+                     num_workers)
+
+    if flow == 'source_to_target':
+        row, col = edge_index[1], edge_index[0]
+    else:
+        row, col = edge_index[0], edge_index[1]
+
+    if not loop:
+        mask = row != col
+        row, col = row[mask], col[mask]
+
+    return torch.stack([row, col], dim=0)

torch_cluster/nearest.py

+from typing import Optional
+
+import torch
+import scipy.cluster
+
+
+def nearest(x: torch.Tensor, y: torch.Tensor,
+            batch_x: Optional[torch.Tensor] = None,
+            batch_y: Optional[torch.Tensor] = None) -> torch.Tensor:
+    r"""Clusters points in :obj:`x` together which are nearest to a given query
+    point in :obj:`y`.
+
+    Args:
+        x (Tensor): Node feature matrix
+            :math:`\mathbf{X} \in \mathbb{R}^{N \times F}`.
+        y (Tensor): Node feature matrix
+            :math:`\mathbf{Y} \in \mathbb{R}^{M \times F}`.
+        batch_x (LongTensor, optional): Batch vector
+            :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each
+            node to a specific example. :obj:`batch_x` needs to be sorted.
+            (default: :obj:`None`)
+        batch_y (LongTensor, optional): Batch vector
+            :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^M`, which assigns each
+            node to a specific example. :obj:`batch_y` needs to be sorted.
+            (default: :obj:`None`)
+
+    :rtype: :class:`LongTensor`
+
+    .. code-block:: python
+
+        import torch
+        from torch_cluster import nearest
+
+        x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]])
+        batch_x = torch.tensor([0, 0, 0, 0])
+        y = torch.Tensor([[-1, 0], [1, 0]])
+        batch_y = torch.tensor([0, 0])
+        cluster = nearest(x, y, batch_x, batch_y)
+    """
+
+    x = x.view(-1, 1) if x.dim() == 1 else x
+    y = y.view(-1, 1) if y.dim() == 1 else y
+    assert x.size(1) == y.size(1)
+
+    if x.is_cuda:
+        if batch_x is not None:
+            assert x.size(0) == batch_x.numel()
+            batch_size = int(batch_x.max()) + 1
+
+            deg = x.new_zeros(batch_size, dtype=torch.long)
+            deg.scatter_add_(0, batch_x, torch.ones_like(batch_x))
+
+            ptr_x = deg.new_zeros(batch_size + 1)
+            torch.cumsum(deg, 0, out=ptr_x[1:])
+        else:
+            ptr_x = torch.tensor([0, x.size(0)], device=x.device)
+
+        if batch_y is not None:
+            assert y.size(0) == batch_y.numel()
+            batch_size = int(batch_y.max()) + 1
+
+            deg = y.new_zeros(batch_size, dtype=torch.long)
+            deg.scatter_add_(0, batch_y, torch.ones_like(batch_y))
+
+            ptr_y = deg.new_zeros(batch_size + 1)
+            torch.cumsum(deg, 0, out=ptr_y[1:])
+        else:
+            ptr_y = torch.tensor([0, y.size(0)], device=y.device)
+
+        return torch.ops.torch_cluster.nearest(x, y, ptr_x, ptr_y)
+    else:
+        # Translate and rescale x and y to [0, 1].
+        if batch_x is not None and batch_y is not None:
+            assert x.dim() == 2 and batch_x.dim() == 1
+            assert y.dim() == 2 and batch_y.dim() == 1
+            assert x.size(0) == batch_x.size(0)
+            assert y.size(0) == batch_y.size(0)
+
+            min_xy = min(x.min().item(), y.min().item())
+            x, y = x - min_xy, y - min_xy
+
+            max_xy = max(x.max().item(), y.max().item())
+            x.div_(max_xy)
+            y.div_(max_xy)
+
+            # Concat batch/features to ensure no cross-links between examples.
+            D = x.size(-1)
+            x = torch.cat([x, 2 * D * batch_x.view(-1, 1).to(x.dtype)], -1)
+            y = torch.cat([y, 2 * D * batch_y.view(-1, 1).to(y.dtype)], -1)
+
+        return torch.from_numpy(
+            scipy.cluster.vq.vq(x.detach().cpu(),
+                                y.detach().cpu())[0]).to(torch.long)

torch_cluster/radius.py

+from typing import Optional
+
+import torch
+
+
+@torch.jit.script
+def radius(x: torch.Tensor, y: torch.Tensor, r: float,
+           batch_x: Optional[torch.Tensor] = None,
+           batch_y: Optional[torch.Tensor] = None, max_num_neighbors: int = 32,
+           num_workers: int = 1) -> torch.Tensor:
+    r"""Finds for each element in :obj:`y` all points in :obj:`x` within
+    distance :obj:`r`.
+
+    Args:
+        x (Tensor): Node feature matrix
+            :math:`\mathbf{X} \in \mathbb{R}^{N \times F}`.
+        y (Tensor): Node feature matrix
+            :math:`\mathbf{Y} \in \mathbb{R}^{M \times F}`.
+        r (float): The radius.
+        batch_x (LongTensor, optional): Batch vector
+            :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each
+            node to a specific example. :obj:`batch_x` needs to be sorted.
+            (default: :obj:`None`)
+        batch_y (LongTensor, optional): Batch vector
+            :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^M`, which assigns each
+            node to a specific example. :obj:`batch_y` needs to be sorted.
+            (default: :obj:`None`)
+        max_num_neighbors (int, optional): The maximum number of neighbors to
+            return for each element in :obj:`y`.
+            If the number of actual neighbors is greater than
+            :obj:`max_num_neighbors`, returned neighbors are picked randomly.
+            (default: :obj:`32`)
+        num_workers (int): Number of workers to use for computation. Has no
+            effect in case :obj:`batch_x` or :obj:`batch_y` is not
+            :obj:`None`, or the input lies on the GPU. (default: :obj:`1`)
+
+    .. code-block:: python
+
+        import torch
+        from torch_cluster import radius
+
+        x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]])
+        batch_x = torch.tensor([0, 0, 0, 0])
+        y = torch.Tensor([[-1, 0], [1, 0]])
+        batch_y = torch.tensor([0, 0])
+        assign_index = radius(x, y, 1.5, batch_x, batch_y)
+    """
+
+    x = x.view(-1, 1) if x.dim() == 1 else x
+    y = y.view(-1, 1) if y.dim() == 1 else y
+    x, y = x.contiguous(), y.contiguous()
+
+    batch_size = 1
+    if batch_x is not None:
+        assert x.size(0) == batch_x.numel()
+        batch_size = int(batch_x.max()) + 1
+    if batch_y is not None:
+        assert y.size(0) == batch_y.numel()
+        batch_size = max(batch_size, int(batch_y.max()) + 1)
+
+    ptr_x: Optional[torch.Tensor] = None
+    ptr_y: Optional[torch.Tensor] = None
+    if batch_size > 1:
+        assert batch_x is not None
+        assert batch_y is not None
+        arange = torch.arange(batch_size + 1, device=x.device)
+        ptr_x = torch.bucketize(arange, batch_x)
+        ptr_y = torch.bucketize(arange, batch_y)
+
+    return torch.ops.torch_cluster.radius(x, y, ptr_x, ptr_y, r,
+                                          max_num_neighbors, num_workers)
+
+
+@torch.jit.script
+def radius_graph(x: torch.Tensor, r: float,
+                 batch: Optional[torch.Tensor] = None, loop: bool = False,
+                 max_num_neighbors: int = 32, flow: str = 'source_to_target',
+                 num_workers: int = 1) -> torch.Tensor:
+    r"""Computes graph edges to all points within a given distance.
+
+    Args:
+        x (Tensor): Node feature matrix
+            :math:`\mathbf{X} \in \mathbb{R}^{N \times F}`.
+        r (float): The radius.
+        batch (LongTensor, optional): Batch vector
+            :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each
+            node to a specific example. :obj:`batch` needs to be sorted.
+            (default: :obj:`None`)
+        loop (bool, optional): If :obj:`True`, the graph will contain
+            self-loops. (default: :obj:`False`)
+        max_num_neighbors (int, optional): The maximum number of neighbors to
+            return for each element.
+            If the number of actual neighbors is greater than
+            :obj:`max_num_neighbors`, returned neighbors are picked randomly.
+            (default: :obj:`32`)
+        flow (string, optional): The flow direction when used in combination
+            with message passing (:obj:`"source_to_target"` or
+            :obj:`"target_to_source"`). (default: :obj:`"source_to_target"`)
+        num_workers (int): Number of workers to use for computation. Has no
+            effect in case :obj:`batch` is not :obj:`None`, or the input lies
+            on the GPU. (default: :obj:`1`)
+
+    :rtype: :class:`LongTensor`
+
+    .. code-block:: python
+
+        import torch
+        from torch_cluster import radius_graph
+
+        x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]])
+        batch = torch.tensor([0, 0, 0, 0])
+        edge_index = radius_graph(x, r=1.5, batch=batch, loop=False)
+    """
+
+    assert flow in ['source_to_target', 'target_to_source']
+    edge_index = radius(x, x, r, batch, batch,
+                        max_num_neighbors if loop else max_num_neighbors + 1,
+                        num_workers)
+    if flow == 'source_to_target':
+        row, col = edge_index[1], edge_index[0]
+    else:
+        row, col = edge_index[0], edge_index[1]
+
+    if not loop:
+        mask = row != col
+        row, col = row[mask], col[mask]
+
+    return torch.stack([row, col], dim=0)

torch_cluster/rw.py

+from typing import Optional
+
+import torch
+from torch import Tensor
+
+
+@torch.jit.script
+def random_walk(row: Tensor, col: Tensor, start: Tensor, walk_length: int,
+                p: float = 1, q: float = 1, coalesced: bool = True,
+                num_nodes: Optional[int] = None) -> Tensor:
+    """Samples random walks of length :obj:`walk_length` from all node indices
+    in :obj:`start` in the graph given by :obj:`(row, col)` as described in the
+    `"node2vec: Scalable Feature Learning for Networks"
+    <https://arxiv.org/abs/1607.00653>`_ paper.
+    Edge indices :obj:`(row, col)` need to be coalesced/sorted according to
+    :obj:`row` (use the :obj:`coalesced` attribute to force).
+
+    Args:
+        row (LongTensor): Source nodes.
+        col (LongTensor): Target nodes.
+        start (LongTensor): Nodes from where random walks start.
+        walk_length (int): The walk length.
+        p (float, optional): Likelihood of immediately revisiting a node in the
+            walk. (default: :obj:`1`)
+        q (float, optional): Control parameter to interpolate between
+            breadth-first strategy and depth-first strategy (default: :obj:`1`)
+        coalesced (bool, optional): If set to :obj:`True`, will coalesce/sort
+            the graph given by :obj:`(row, col)` according to :obj:`row`.
+            (default: :obj:`True`)
+        num_nodes (int, optional): The number of nodes. (default: :obj:`None`)
+
+    :rtype: :class:`LongTensor`
+    """
+    if num_nodes is None:
+        num_nodes = max(int(row.max()), int(col.max()), int(start.max())) + 1
+
+    if coalesced:
+        perm = torch.argsort(row * num_nodes + col)
+        row, col = row[perm], col[perm]
+
+    deg = row.new_zeros(num_nodes)
+    deg.scatter_add_(0, row, torch.ones_like(row))
+    rowptr = row.new_zeros(num_nodes + 1)
+    torch.cumsum(deg, 0, out=rowptr[1:])
+
+    return torch.ops.torch_cluster.random_walk(rowptr, col, start, walk_length,
+                                               p, q)[0]

torch_cluster/sampler.py

+import torch
+
+
+@torch.jit.script
+def neighbor_sampler(start: torch.Tensor, rowptr: torch.Tensor, size: float):
+    assert not start.is_cuda
+
+    factor: float = -1.
+    count: int = -1
+    if size <= 1:
+        factor = size
+        assert factor > 0
+    else:
+        count = int(size)
+
+    return torch.ops.torch_cluster.neighbor_sampler(start, rowptr, count,
+                                                    factor)