year up, restricted coverage, nested extensions

rusty1s · rusty1s · commit dfe188ab6e91 · 2019-01-24T10:33:03.000+01:00
diff --git a/.coveragerc b/.coveragerc
@@ -1,3 +1,5 @@
+[run]
+source=torch_cluster
 [report]
 exclude_lines =
     pragma: no cover
diff --git a/LICENSE b/LICENSE
@@ -1,4 +1,4 @@
-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
diff --git a/cuda/fps.cpp b/cuda/fps.cpp
@@ -1,4 +1,4 @@
-#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");
diff --git a/cuda/graclus.cpp b/cuda/graclus.cpp
@@ -1,4 +1,4 @@
-#include <torch/torch.h>
+#include <torch/extension.h>
 
 #define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be CUDA tensor")
 
diff --git a/cuda/grid.cpp b/cuda/grid.cpp
@@ -1,4 +1,4 @@
-#include <torch/torch.h>
+#include <torch/extension.h>
 
 #define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be CUDA tensor")
 
diff --git a/cuda/knn.cpp b/cuda/knn.cpp
@@ -1,4 +1,4 @@
-#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");
diff --git a/cuda/nearest.cpp b/cuda/nearest.cpp
@@ -1,4 +1,4 @@
-#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");
diff --git a/cuda/radius.cpp b/cuda/radius.cpp
@@ -1,4 +1,4 @@
-#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");
diff --git a/setup.py b/setup.py
@@ -3,26 +3,29 @@
 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 @@
 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 @@
     tests_require=tests_require,
     ext_modules=ext_modules,
     cmdclass=cmdclass,
-    packages=find_packages(), )
+    packages=find_packages(),
+)
diff --git a/torch_cluster/__init__.py b/torch_cluster/__init__.py
@@ -5,7 +5,7 @@
 from .knn import knn, knn_graph
 from .radius import radius, radius_graph
 
-__version__ = '1.2.2'
+__version__ = '1.2.3'
 
 __all__ = [
     'graclus_cluster',
diff --git a/torch_cluster/fps.py b/torch_cluster/fps.py
@@ -1,8 +1,8 @@
 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)
diff --git a/torch_cluster/graclus.py b/torch_cluster/graclus.py
@@ -1,8 +1,8 @@
 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)
diff --git a/torch_cluster/grid.py b/torch_cluster/grid.py
@@ -1,8 +1,8 @@
 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
diff --git a/torch_cluster/knn.py b/torch_cluster/knn.py
@@ -2,7 +2,7 @@
 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())
diff --git a/torch_cluster/nearest.py b/torch_cluster/nearest.py
@@ -2,7 +2,7 @@
 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())
diff --git a/torch_cluster/radius.py b/torch_cluster/radius.py
@@ -2,7 +2,7 @@
 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)

Original file line number	Diff line number	Diff line change
`@@ -1,3 +1,5 @@`
	`1`	`+[run]`
	`2`	`+source=torch_cluster`
`1`	`3`	`[report]`
`2`	`4`	`exclude_lines =`
`3`	`5`	`pragma: no cover`
Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-Copyright (c) 2018 Matthias Fey <matthias.fey@tu-dortmund.de>`
	`1`	`+Copyright (c) 2019 Matthias Fey <matthias.fey@tu-dortmund.de>`
`2`	`2`
`3`	`3`	`Permission is hereby granted, free of charge, to any person obtaining a copy`
`4`	`4`	`of this software and associated documentation files (the "Software"), to deal`
Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-#include <torch/torch.h>`
	`1`	`+#include <torch/extension.h>`
`2`	`2`
`3`	`3`	`#define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be CUDA tensor")`
`4`	`4`	`#define IS_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " is not contiguous");`