2 files changed, 11 insertions, 4 deletions
diff --git a/third_party/mpi/BUILD b/third_party/mpi/BUILD
index ec09b045b7..ff3f437e92 100644
--- a/third_party/mpi/BUILD
+++ b/third_party/mpi/BUILD
@@ -13,12 +13,13 @@ filegroup(
 )
 
 load("//third_party/mpi:mpi.bzl", "mpi_hdr")
+load("//third_party/mpi:mpi.bzl", "if_mpi")
 
 cc_library(
     name = "mpi",
-    srcs = select({
-        "//conditions:default": ["libmpi.so"],
-    }),
-    hdrs = mpi_hdr(),
+    srcs = if_mpi([
+        "libmpi.so",
+    ]),
+    hdrs = if_mpi(mpi_hdr()),
     visibility = ["//visibility:public"],
 )
diff --git a/third_party/mpi/mpi.bzl b/third_party/mpi/mpi.bzl
index 1f7d2f4102..38ce91c4d0 100644
--- a/third_party/mpi/mpi.bzl
+++ b/third_party/mpi/mpi.bzl
@@ -9,3 +9,9 @@ def mpi_hdr():
     else:
         hdrs = ["mpi.h",  "mpio.h", "mpicxx.h"]        #When using MVAPICH
     return hdrs
+
+def if_mpi(if_true, if_false = []):
+    return select({
+        "//tensorflow:with_mpi_support": if_true,
+        "//conditions:default": if_false
+    })