update

setup.py

@@ -26,7 +26,6 @@ def get_extensions():
         define_macros += [('WITH_CUDA', None)]
         nvcc_flags = os.getenv('NVCC_FLAGS', '')
         nvcc_flags = [] if nvcc_flags == '' else nvcc_flags.split(' ')
-        nvcc_flags += ['-arch=sm_35', '--expt-relaxed-constexpr']
         extra_compile_args['nvcc'] = nvcc_flags
 
     extensions_dir = osp.join(osp.dirname(osp.abspath(__file__)), 'csrc')
@@ -76,6 +75,7 @@ setup(
         'cluster-algorithms',
     ],
     license='MIT',
+    python_requires='>=3.6',
     install_requires=install_requires,
     setup_requires=setup_requires,
     tests_require=tests_require,

test/utils.py

 import torch
-from torch.testing import get_all_dtypes
-
-dtypes = get_all_dtypes()
-dtypes.remove(torch.half)
-dtypes.remove(torch.bool)
-if hasattr(torch, 'bfloat16'):
-    dtypes.remove(torch.bfloat16)
 
+dtypes = [torch.float, torch.double, torch.int, torch.long]
 grad_dtypes = [torch.float, torch.double]
 
 devices = [torch.device('cpu')]
 if torch.cuda.is_available():
-    devices += [torch.device('cuda:{}'.format(torch.cuda.current_device()))]
+    devices += [torch.device(f'cuda:{torch.cuda.current_device()}')]
 
 
 def tensor(x, dtype, device):

torch_cluster/__init__.py

-from .graclus import graclus_cluster
-from .grid import grid_cluster
-from .fps import fps
-from .nearest import nearest
-from .knn import knn, knn_graph
-from .radius import radius, radius_graph
-from .rw import random_walk
-from .sampler import neighbor_sampler
+import importlib
+import os.path as osp
+
+import torch
 
 __version__ = '1.5.0'
+expected_torch_version = (1, 4)
+
+try:
+    for library in ['_version', '_grid', '_graclus', '_fps']:
+        torch.ops.load_library(importlib.machinery.PathFinder().find_spec(
+            library, [osp.dirname(__file__)]).origin)
+except OSError as e:
+    major, minor = [int(x) for x in torch.__version__.split('.')[:2]]
+    t_major, t_minor = expected_torch_version
+    if major != t_major or (major == t_major and minor != t_minor):
+        raise RuntimeError(
+            f'Expected PyTorch version {t_major}.{t_minor} but found '
+            f'version {major}.{minor}.')
+    raise OSError(e)
+
+if torch.version.cuda is not None:  # pragma: no cover
+    cuda_version = torch.ops.torch_sparse.cuda_version()
+
+    if cuda_version == -1:
+        major = minor = 0
+    elif 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])
+    t_major, t_minor = [int(x) for x in torch.version.cuda.split('.')]
+
+    if t_major != major or t_minor != minor:
+        raise RuntimeError(
+            f'Detected that PyTorch and torch_cluster were compiled with '
+            f'different CUDA versions. PyTorch has CUDA version '
+            f'{t_major}.{t_minor} and torch_cluster has CUDA version '
+            f'{major}.{minor}. Please reinstall the torch_cluster that '
+            f'matches your PyTorch install.')
+
+from .graclus import graclus_cluster  # noqa
+from .grid import grid_cluster  # noqa
+from .fps import fps  # noqa
+# from .nearest import nearest  # noqa
+# from .knn import knn, knn_graph  # 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',
+    # 'nearest',
+    # 'knn',
+    # 'knn_graph',
+    # 'radius',
+    # 'radius_graph',
+    # 'random_walk',
+    # 'neighbor_sampler',
     '__version__',
 ]

torch_cluster/fps.py

-import torch
-import torch_cluster.fps_cpu
+from typing import Optional
 
-if torch.cuda.is_available():
-    import torch_cluster.fps_cuda
+import torch
 
 
-def fps(x, batch=None, ratio=0.5, random_start=True):
+@torch.jit.script
+def fps(src: torch.Tensor, batch: Optional[torch.Tensor] = None,
+        ratio: float = 0.5, random_start: bool = True) -> torch.Tensor:
     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:
-        x (Tensor): Node feature matrix
+        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
@@ -23,28 +23,25 @@ def fps(x, batch=None, ratio=0.5, random_start=True):
 
     :rtype: :class:`LongTensor`
 
-    .. testsetup::
+    .. code-block:: python
 
         import torch
         from torch_cluster import fps
 
-    .. testcode::
-
-        >>> x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]])
-        >>> batch = torch.tensor([0, 0, 0, 0])
-        >>> index = fps(x, batch, ratio=0.5)
+        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)
     """
 
-    if batch is None:
-        batch = x.new_zeros(x.size(0), dtype=torch.long)
+    ptr: Optional[torch.Tensor] = None
+    if batch is not None:
+        assert src.size(0) == batch.size(0)
+        batch_size = int(batch.max()) + 1
 
-    x = x.view(-1, 1) if x.dim() == 1 else x
+        deg = src.new_zeros(batch_size, dtype=torch.long)
+        deg.scatter_add_(0, batch, torch.ones_like(batch))
 
-    assert x.dim() == 2 and batch.dim() == 1
-    assert x.size(0) == batch.size(0)
-    assert ratio > 0 and ratio < 1
+        ptr = src.new_zeros(batch_size + 1, dtype=torch.long)
+        deg.cumsum(0, out=ptr[1:])
 
-    if x.is_cuda:
-        return torch_cluster.fps_cuda.fps(x, batch, ratio, random_start)
-    else:
-        return torch_cluster.fps_cpu.fps(x, batch, ratio, random_start)
+    return torch.ops.torch_cluster.fps(src, ptr, ratio, random_start)

torch_cluster/graclus.py

-import torch
-import torch_cluster.graclus_cpu
+from typing import Optional
 
-if torch.cuda.is_available():
-    import torch_cluster.graclus_cuda
+import torch
 
 
-def graclus_cluster(row, col, weight=None, num_nodes=None):
+@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).
 
@@ -17,25 +18,18 @@ def graclus_cluster(row, col, weight=None, num_nodes=None):
 
     :rtype: :class:`LongTensor`
 
-    Examples::
+    .. 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)
+        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(row.max().item(), col.max().item()) + 1
-
-    if row.is_cuda:
-        op = torch_cluster.graclus_cuda
-    else:
-        op = torch_cluster.graclus_cpu
-
-    if weight is None:
-        cluster = op.graclus(row, col, num_nodes)
-    else:
-        cluster = op.weighted_graclus(row, col, weight, num_nodes)
+        num_nodes = max(int(row.max()), int(col.max())) + 1
 
-    return cluster
+    return torch.ops.torch_cluster.graclus(row, col, weight, num_nodes)

torch_cluster/grid.py

-import torch
-import torch_cluster.grid_cpu
+from typing import Optional
 
-if torch.cuda.is_available():
-    import torch_cluster.grid_cuda
+import torch
 
 
-def grid_cluster(pos, size, start=None, end=None):
+@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.
 
@@ -19,22 +20,13 @@ def grid_cluster(pos, size, start=None, end=None):
 
     :rtype: :class:`LongTensor`
 
-    Examples::
-
-        >>> pos = torch.Tensor([[0, 0], [11, 9], [2, 8], [2, 2], [8, 3]])
-        >>> size = torch.Tensor([5, 5])
-        >>> cluster = grid_cluster(pos, size)
-    """
+    .. code-block:: python
 
-    pos = pos.unsqueeze(-1) if pos.dim() == 1 else pos
-    start = pos.t().min(dim=1)[0] if start is None else start
-    end = pos.t().max(dim=1)[0] if end is None else end
+        import torch
+        from torch_cluster import grid_cluster
 
-    if pos.is_cuda:
-        op = torch_cluster.grid_cuda
-    else:
-        op = torch_cluster.grid_cpu
-
-    cluster = op.grid(pos, size, start, end)
-
-    return 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)