WITH_PYTHON conditionals (#313)

Author: GericoVi

CMakeLists.txt

@@ -4,6 +4,7 @@ set(CMAKE_CXX_STANDARD 14)
 set(TORCHSCATTER_VERSION 2.0.9)
 
 option(WITH_CUDA "Enable CUDA support" OFF)
+option(WITH_PYTHON "Link to Python when building" ON)
 
 if(WITH_CUDA)
   enable_language(CUDA)
@@ -12,7 +13,10 @@ if(WITH_CUDA)
   set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --expt-relaxed-constexpr")
 endif()
 
-find_package(Python3 COMPONENTS Development)
+if (WITH_PYTHON)
+  add_definitions(-DWITH_PYTHON)
+  find_package(Python3 COMPONENTS Development)
+endif()
 find_package(Torch REQUIRED)
 
 file(GLOB HEADERS csrc/*.h)
@@ -22,7 +26,10 @@ if(WITH_CUDA)
 endif()
 
 add_library(${PROJECT_NAME} SHARED ${OPERATOR_SOURCES})
-target_link_libraries(${PROJECT_NAME} PRIVATE ${TORCH_LIBRARIES} Python3::Python)
+target_link_libraries(${PROJECT_NAME} PRIVATE ${TORCH_LIBRARIES})
+if (WITH_PYTHON)
+  target_link_libraries(${PROJECT_NAME} PRIVATE Python3::Python)
+endif()
 set_target_properties(${PROJECT_NAME} PROPERTIES EXPORT_NAME TorchScatter)
 
 target_include_directories(${PROJECT_NAME} INTERFACE

csrc/cpu/index_info.h

@@ -1,6 +1,6 @@
 #pragma once
 
-#include <torch/extension.h>
+#include "../extensions.h"
 
 #define MAX_TENSORINFO_DIMS 25
 

csrc/cpu/scatter_cpu.h

@@ -1,6 +1,6 @@
 #pragma once
 
-#include <torch/extension.h>
+#include "../extensions.h"
 
 std::tuple<torch::Tensor, torch::optional<torch::Tensor>>
 scatter_cpu(torch::Tensor src, torch::Tensor index, int64_t dim,

csrc/cpu/segment_coo_cpu.h

@@ -1,6 +1,6 @@
 #pragma once
 
-#include <torch/extension.h>
+#include "../extensions.h"
 
 std::tuple<torch::Tensor, torch::optional<torch::Tensor>>
 segment_coo_cpu(torch::Tensor src, torch::Tensor index,

csrc/cpu/segment_csr_cpu.h

@@ -1,6 +1,6 @@
 #pragma once
 
-#include <torch/extension.h>
+#include "../extensions.h"
 
 std::tuple<torch::Tensor, torch::optional<torch::Tensor>>
 segment_csr_cpu(torch::Tensor src, torch::Tensor indptr,

csrc/cpu/utils.h

@@ -1,6 +1,6 @@
 #pragma once
 
-#include <torch/extension.h>
+#include "../extensions.h"
 
 #define CHECK_CPU(x) AT_ASSERTM(x.device().is_cpu(), #x " must be CPU tensor")
 #define CHECK_INPUT(x) AT_ASSERTM(x, "Input mismatch")

csrc/cuda/scatter_cuda.h

@@ -1,6 +1,6 @@
 #pragma once
 
-#include <torch/extension.h>
+#include "../extensions.h"
 
 std::tuple<torch::Tensor, torch::optional<torch::Tensor>>
 scatter_cuda(torch::Tensor src, torch::Tensor index, int64_t dim,

csrc/cuda/segment_coo_cuda.h

@@ -1,6 +1,6 @@
 #pragma once
 
-#include <torch/extension.h>
+#include "../extensions.h"
 
 std::tuple<torch::Tensor, torch::optional<torch::Tensor>>
 segment_coo_cuda(torch::Tensor src, torch::Tensor index,

csrc/cuda/segment_csr_cuda.h

@@ -1,6 +1,6 @@
 #pragma once
 
-#include <torch/extension.h>
+#include "../extensions.h"
 
 std::tuple<torch::Tensor, torch::optional<torch::Tensor>>
 segment_csr_cuda(torch::Tensor src, torch::Tensor indptr,

csrc/cuda/utils.cuh

@@ -1,6 +1,6 @@
 #pragma once
 
-#include <torch/extension.h>
+#include "../extensions.h"
 
 #define CHECK_CUDA(x)                                                          \
   AT_ASSERTM(x.device().is_cuda(), #x " must be CUDA tensor")