[mlgo] Fix cmake logic detecting tf pip package location

New logic works for both `tensorflow` and `tf-nightly`.
This commit is contained in:
Mircea Trofin 2022-08-24 08:28:23 -07:00
parent 80cc8f0f62
commit bbf11c8d9f
1 changed files with 4 additions and 6 deletions

View File

@ -962,15 +962,13 @@ if (TENSORFLOW_C_LIB_PATH)
if (NOT TF_PROTO_HEADERS)
message(STATUS "TF_PROTO_HEADERS not defined. Looking for tensorflow pip package.")
execute_process(COMMAND
${Python3_EXECUTABLE} "-m" "pip" "show" "tensorflow"
OUTPUT_VARIABLE TF_PIP_OUT)
if ("${TF_PIP_OUT}" STREQUAL "")
${Python3_EXECUTABLE} "-c" "import tensorflow as tf; import os; print(os.path.dirname(tf.__file__))"
OUTPUT_VARIABLE TF_PIP)
if ("${TF_PIP}" STREQUAL "")
message(FATAL ERROR "Tensorflow pip package is also required for 'development' mode (protobuf headers)")
endif()
string(REGEX MATCH "Location: ([^\n]*\n)" TF_PIP_LOC "${TF_PIP_OUT}")
string(REPLACE "Location: " "" TF_PIP ${TF_PIP_LOC})
string(STRIP ${TF_PIP} TF_PIP)
set(TF_PROTO_HEADERS "${TF_PIP}/tensorflow/include")
set(TF_PROTO_HEADERS "${TF_PIP}/include")
endif()
message(STATUS "Using Tensorflow headers under: ${TF_PROTO_HEADERS}")
include_directories(${TF_PROTO_HEADERS})