diff --git a/llvm/lib/Analysis/CMakeLists.txt b/llvm/lib/Analysis/CMakeLists.txt
--- a/llvm/lib/Analysis/CMakeLists.txt
+++ b/llvm/lib/Analysis/CMakeLists.txt
@@ -1,4 +1,4 @@
-if (DEFINED LLVM_HAVE_TF_AOT OR DEFINED LLVM_HAVE_TF_API)
+if (DEFINED LLVM_HAVE_TF_AOT OR LLVM_HAVE_TFLITE)
   include(TensorFlowCompile)
   set(LLVM_INLINER_MODEL_PATH_DEFAULT "models/inliner-Oz")
 
@@ -17,18 +17,14 @@
     )
   endif()
 
-  if (DEFINED LLVM_HAVE_TF_API)
-    if (LLVM_HAVE_TFLITE)
-      build_proto(
-        tensorflow/core/protobuf/error_codes
-        tensorflow/core/example/feature
-        tensorflow/core/example/example)
-      list(APPEND MLLinkDeps 
-        tensorflow-lite::tensorflow-lite
-        protobuf::libprotobuf)
-    else()
-      list(APPEND MLLinkDeps ${tensorflow_c_api} ${tensorflow_fx})
-    endif()
+  if (LLVM_HAVE_TFLITE)
+    build_proto(
+      tensorflow/core/protobuf/error_codes
+      tensorflow/core/example/feature
+      tensorflow/core/example/example)
+    list(APPEND MLLinkDeps
+      tensorflow-lite::tensorflow-lite
+      protobuf::libprotobuf)
   endif()
 endif()
 
diff --git a/llvm/unittests/Analysis/CMakeLists.txt b/llvm/unittests/Analysis/CMakeLists.txt
--- a/llvm/unittests/Analysis/CMakeLists.txt
+++ b/llvm/unittests/Analysis/CMakeLists.txt
@@ -7,7 +7,7 @@
   )
 
 set(MLGO_TESTS TFUtilsTest.cpp TrainingLoggerTest.cpp)
-if (DEFINED LLVM_HAVE_TF_API)
+if (LLVM_HAVE_TFLITE)
   LIST(APPEND EXTRA_TESTS ${MLGO_TESTS})
 else()
   LIST(APPEND LLVM_OPTIONAL_SOURCES ${MLGO_TESTS})