nested extensions

Author: rusty1s

setup.py

@@ -4,14 +4,14 @@
 
 ext_modules = [
     CppExtension(
-        'scatter_cpu', ['cpu/scatter.cpp'],
+        'torch_scatter.scatter_cpu', ['cpu/scatter.cpp'],
         extra_compile_args=['-Wno-unused-variable'])
 ]
 cmdclass = {'build_ext': torch.utils.cpp_extension.BuildExtension}
 
 if CUDA_HOME is not None:
     ext_modules += [
-        CUDAExtension('scatter_cuda',
+        CUDAExtension('torch_scatter.scatter_cuda',
                       ['cuda/scatter.cpp', 'cuda/scatter_kernel.cu'])
     ]
 

torch_scatter/utils/ext.py

@@ -1,10 +1,13 @@
 import torch
-import scatter_cpu
+import torch_scatter.scatter_cpu
 
 if torch.cuda.is_available():
-    import scatter_cuda
+    import torch_scatter.scatter_cuda
 
 
 def get_func(name, tensor):
-    module = scatter_cuda if tensor.is_cuda else scatter_cpu
+    if tensor.is_cuda:
+        module = torch_scatter.scatter_cuda
+    else:
+        module = torch_scatter.scatter_cpu
     return getattr(module, name)