added flag to enable tf pybind support

CMakeLists.txt

@@ -39,6 +39,8 @@ else()
     set(MIGRAPHX_ENABLE_GPU Off CACHE BOOL "")
 endif()
 
+set(MIGRAPHX_ENABLE_TF Off CACHE BOOL "")
+
 add_compile_options(-std=c++14)
 
 list(APPEND CMAKE_MODULE_PATH ${CMAKE_CURRENT_SOURCE_DIR}/cmake)

src/py/CMakeLists.txt

@@ -12,7 +12,12 @@ if(MIGRAPHX_ENABLE_PYTHON)
         C_VISIBILITY_PRESET hidden
         CXX_VISIBILITY_PRESET hidden
     )
-    target_link_libraries(migraphx_py PRIVATE migraphx migraphx_tf migraphx_cpu)
+    if(MIGRAPHX_ENABLE_TF)
+        target_link_libraries(migraphx_py PRIVATE migraphx migraphx_tf migraphx_cpu)
+        target_compile_definitions(migraphx_py PRIVATE -DENABLE_TF)
+    else()
+        target_link_libraries(migraphx_py PRIVATE migraphx migraphx_onnx migraphx_cpu)
+    endif()
     if(MIGRAPHX_ENABLE_GPU)
         target_link_libraries(migraphx_py PRIVATE migraphx_gpu)
         target_compile_definitions(migraphx_py PRIVATE -DHAVE_GPU)

src/py/migraphx_py.cpp

@@ -4,9 +4,13 @@
 #include <migraphx/program.hpp>
 #include <migraphx/generate.hpp>
 #include <migraphx/cpu/target.hpp>
-// #include <migraphx/onnx.hpp>
-#include <migraphx/tf.hpp>
 #include <migraphx/stringutils.hpp>
+#ifdef ENABLE_TF
+#include <migraphx/tf.hpp>
+#else
+#include <migraphx/onnx.hpp>
+#endif
+
 #ifdef HAVE_GPU
 #include <migraphx/gpu/target.hpp>
 #include <migraphx/gpu/hip.hpp>
@@ -156,14 +160,16 @@ PYBIND11_MODULE(migraphx, m)
         .def("__ne__", std::not_equal_to<migraphx::program>{})
         .def("__repr__", [](const migraphx::program& p) { return migraphx::to_string(p); });
 
-    //     m.def("parse_onnx", &migraphx::parse_onnx);
-
+#ifdef ENABLE_TF
     m.def("parse_tf",
           &migraphx::parse_tf,
           "Parse tf protobuf (default format is nhwc)",
           py::arg("filename"),
           py::arg("is_nhwc") = true);
+#else
+    m.def("parse_onnx", &migraphx::parse_onnx);
 
+#endif
     m.def("get_target", [](const std::string& name) -> migraphx::target {
         if(name == "cpu")
             return migraphx::cpu::target{};