year up, restricted coverage, nested extensions

.coveragerc

+[run]
+source=torch_cluster
 [report]
 exclude_lines =
     pragma: no cover

LICENSE

-Copyright (c) 2018 Matthias Fey <matthias.fey@tu-dortmund.de>
+Copyright (c) 2019 Matthias Fey <matthias.fey@tu-dortmund.de>
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

cuda/fps.cpp

-#include <torch/torch.h>
+#include <torch/extension.h>
 
 #define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be CUDA tensor")
 #define IS_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " is not contiguous");

cuda/graclus.cpp

-#include <torch/torch.h>
+#include <torch/extension.h>
 
 #define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be CUDA tensor")
 

cuda/grid.cpp

-#include <torch/torch.h>
+#include <torch/extension.h>
 
 #define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be CUDA tensor")
 

cuda/knn.cpp

-#include <torch/torch.h>
+#include <torch/extension.h>
 
 #define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be CUDA tensor")
 #define IS_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " is not contiguous");

cuda/nearest.cpp

-#include <torch/torch.h>
+#include <torch/extension.h>
 
 #define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be CUDA tensor")
 #define IS_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " is not contiguous");

cuda/radius.cpp

-#include <torch/torch.h>
+#include <torch/extension.h>
 
 #define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be CUDA tensor")
 #define IS_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " is not contiguous");

setup.py

@@ -3,26 +3,29 @@ import torch
 from torch.utils.cpp_extension import CppExtension, CUDAExtension, CUDA_HOME
 
 ext_modules = [
-    CppExtension('graclus_cpu', ['cpu/graclus.cpp']),
-    CppExtension('grid_cpu', ['cpu/grid.cpp']),
-    CppExtension('fps_cpu', ['cpu/fps.cpp']),
+    CppExtension('torch_cluster.graclus_cpu', ['cpu/graclus.cpp']),
+    CppExtension('torch_cluster.grid_cpu', ['cpu/grid.cpp']),
+    CppExtension('torch_cluster.fps_cpu', ['cpu/fps.cpp']),
 ]
 cmdclass = {'build_ext': torch.utils.cpp_extension.BuildExtension}
 
 if CUDA_HOME is not None:
     ext_modules += [
-        CUDAExtension('graclus_cuda',
+        CUDAExtension('torch_cluster.graclus_cuda',
                       ['cuda/graclus.cpp', 'cuda/graclus_kernel.cu']),
-        CUDAExtension('grid_cuda', ['cuda/grid.cpp', 'cuda/grid_kernel.cu']),
-        CUDAExtension('fps_cuda', ['cuda/fps.cpp', 'cuda/fps_kernel.cu']),
-        CUDAExtension('nearest_cuda',
+        CUDAExtension('torch_cluster.grid_cuda',
+                      ['cuda/grid.cpp', 'cuda/grid_kernel.cu']),
+        CUDAExtension('torch_cluster.fps_cuda',
+                      ['cuda/fps.cpp', 'cuda/fps_kernel.cu']),
+        CUDAExtension('torch_cluster.nearest_cuda',
                       ['cuda/nearest.cpp', 'cuda/nearest_kernel.cu']),
-        CUDAExtension('knn_cuda', ['cuda/knn.cpp', 'cuda/knn_kernel.cu']),
-        CUDAExtension('radius_cuda',
+        CUDAExtension('torch_cluster.knn_cuda',
+                      ['cuda/knn.cpp', 'cuda/knn_kernel.cu']),
+        CUDAExtension('torch_cluster.radius_cuda',
                       ['cuda/radius.cpp', 'cuda/radius_kernel.cu']),
     ]
 
-__version__ = '1.2.2'
+__version__ = '1.2.3'
 url = 'https://github.com/rusty1s/pytorch_cluster'
 
 install_requires = ['scipy']
@@ -32,8 +35,8 @@ tests_require = ['pytest', 'pytest-cov']
 setup(
     name='torch_cluster',
     version=__version__,
-    description='PyTorch Extension Library of Optimized Graph Cluster '
-    'Algorithms',
+    description=('PyTorch Extension Library of Optimized Graph Cluster '
+                 'Algorithms'),
     author='Matthias Fey',
     author_email='matthias.fey@tu-dortmund.de',
     url=url,
@@ -44,4 +47,5 @@ setup(
     tests_require=tests_require,
     ext_modules=ext_modules,
     cmdclass=cmdclass,
-    packages=find_packages(), )
+    packages=find_packages(),
+)

torch_cluster/__init__.py

@@ -5,7 +5,7 @@ from .nearest import nearest
 from .knn import knn, knn_graph
 from .radius import radius, radius_graph
 
-__version__ = '1.2.2'
+__version__ = '1.2.3'
 
 __all__ = [
     'graclus_cluster',

torch_cluster/fps.py

 import torch
-import fps_cpu
+import torch_cluster.fps_cpu
 
 if torch.cuda.is_available():
-    import fps_cuda
+    import torch_cluster.fps_cuda
 
 
 def fps(x, batch=None, ratio=0.5, random_start=True):
@@ -45,6 +45,6 @@ def fps(x, batch=None, ratio=0.5, random_start=True):
     assert ratio > 0 and ratio < 1
 
     if x.is_cuda:
-        return fps_cuda.fps(x, batch, ratio, random_start)
+        return torch_cluster.fps_cuda.fps(x, batch, ratio, random_start)
     else:
-        return fps_cpu.fps(x, batch, ratio, random_start)
+        return torch_cluster.fps_cpu.fps(x, batch, ratio, random_start)

torch_cluster/graclus.py

 import torch
-import graclus_cpu
+import torch_cluster.graclus_cpu
 
 if torch.cuda.is_available():
-    import graclus_cuda
+    import torch_cluster.graclus_cuda
 
 
 def graclus_cluster(row, col, weight=None, num_nodes=None):
@@ -28,7 +28,10 @@ def graclus_cluster(row, col, weight=None, num_nodes=None):
     if num_nodes is None:
         num_nodes = max(row.max().item(), col.max().item()) + 1
 
-    op = graclus_cuda if row.is_cuda else graclus_cpu
+    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)

torch_cluster/grid.py

 import torch
-import grid_cpu
+import torch_cluster.grid_cpu
 
 if torch.cuda.is_available():
-    import grid_cuda
+    import torch_cluster.grid_cuda
 
 
 def grid_cluster(pos, size, start=None, end=None):
@@ -30,7 +30,11 @@ def grid_cluster(pos, size, start=None, end=None):
     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
 
-    op = grid_cuda if pos.is_cuda else grid_cpu
+    if pos.is_cuda:
+        op = torch_cluster.grid_cuda
+    else:
+        op = torch_cluster.grid_cpu
+
     cluster = op.grid(pos, size, start, end)
 
     return cluster

torch_cluster/knn.py

@@ -2,7 +2,7 @@ import torch
 import scipy.spatial
 
 if torch.cuda.is_available():
-    import knn_cuda
+    import torch_cluster.knn_cuda
 
 
 def knn(x, y, k, batch_x=None, batch_y=None):
@@ -54,7 +54,7 @@ def knn(x, y, k, batch_x=None, batch_y=None):
     assert y.size(0) == batch_y.size(0)
 
     if x.is_cuda:
-        return knn_cuda.knn(x, y, k, batch_x, batch_y)
+        return torch_cluster.knn_cuda.knn(x, y, k, batch_x, batch_y)
 
     # Rescale x and y.
     min_xy = min(x.min().item(), y.min().item())

torch_cluster/nearest.py

@@ -2,7 +2,7 @@ import torch
 import scipy.cluster
 
 if torch.cuda.is_available():
-    import nearest_cuda
+    import torch_cluster.nearest_cuda
 
 
 def nearest(x, y, batch_x=None, batch_y=None):
@@ -51,7 +51,7 @@ def nearest(x, y, batch_x=None, batch_y=None):
     assert y.size(0) == batch_y.size(0)
 
     if x.is_cuda:
-        return nearest_cuda.nearest(x, y, batch_x, batch_y)
+        return torch_cluster.nearest_cuda.nearest(x, y, batch_x, batch_y)
 
     # Rescale x and y.
     min_xy = min(x.min().item(), y.min().item())

torch_cluster/radius.py

@@ -2,7 +2,7 @@ import torch
 import scipy.spatial
 
 if torch.cuda.is_available():
-    import radius_cuda
+    import torch_cluster.radius_cuda
 
 
 def radius(x, y, r, batch_x=None, batch_y=None, max_num_neighbors=32):
@@ -57,7 +57,8 @@ def radius(x, y, r, batch_x=None, batch_y=None, max_num_neighbors=32):
     assert y.size(0) == batch_y.size(0)
 
     if x.is_cuda:
-        return radius_cuda.radius(x, y, r, batch_x, batch_y, max_num_neighbors)
+        return torch_cluster.radius_cuda.radius(x, y, r, batch_x, batch_y,
+                                                max_num_neighbors)
 
     x = torch.cat([x, 2 * r * batch_x.view(-1, 1).to(x.dtype)], dim=-1)
     y = torch.cat([y, 2 * r * batch_y.view(-1, 1).to(y.dtype)], dim=-1)