diff --git a/WORKSPACE b/WORKSPACE
index c8bd6bc79007..08aaca272203 100644
--- a/WORKSPACE
+++ b/WORKSPACE
@@ -13,7 +13,6 @@ http_archive(
     ],
     patch_tool = "patch",
     patches = [
-        "//tf_patches:bazel.diff",
         "//tf_patches:cache_urls.diff",
         "//tf_patches:cudnn_int8x32.diff",
         "//tf_patches:f16_abi_clang.diff",
diff --git a/bazel/BUILD b/bazel/BUILD
new file mode 100644
index 000000000000..e69de29bb2d1
diff --git a/bazel/tensorflow.bzl b/bazel/tensorflow.bzl
new file mode 100644
index 000000000000..c5444f45a067
--- /dev/null
+++ b/bazel/tensorflow.bzl
@@ -0,0 +1,8 @@
+"""Macros for working with tensorflow deps."""
+
+def if_with_tpu_support(if_true, if_false = []):
+    """Shorthand for select()ing whether to build API support for TPUs when building TensorFlow"""
+    return select({
+        "@org_tensorflow//tensorflow:with_tpu_support": if_true,
+        "//conditions:default": if_false,
+    })
diff --git a/docker/experimental/Dockerfile b/docker/experimental/Dockerfile
index c236cc8c4914..de23e9e5bfe1 100644
--- a/docker/experimental/Dockerfile
+++ b/docker/experimental/Dockerfile
@@ -73,6 +73,7 @@ COPY tf_patches/ tf_patches/
 COPY .bazelrc .
 COPY .bazelversion .
 COPY WORKSPACE .
+COPY bazel/ bazel/
 COPY build_torch_xla_libs.sh .
 
 COPY torch_xla/ torch_xla/
diff --git a/tf_patches/bazel.diff b/tf_patches/bazel.diff
deleted file mode 100644
index 4b670953bcb1..000000000000
--- a/tf_patches/bazel.diff
+++ /dev/null
@@ -1,13 +0,0 @@
-diff --git i/tensorflow/tensorflow.bzl w/tensorflow/tensorflow.bzl
-index 649c8e22dcc..a85f4bc3af3 100644
---- i/tensorflow/tensorflow.bzl
-+++ w/tensorflow/tensorflow.bzl
-@@ -315,7 +315,7 @@ def if_libtpu(if_true, if_false = []):
- def if_with_tpu_support(if_true, if_false = []):
-     """Shorthand for select()ing whether to build API support for TPUs when building TensorFlow"""
-     return select({
--        "//tensorflow:with_tpu_support": if_true,
-+        clean_dep("//tensorflow:with_tpu_support"): if_true,
-         "//conditions:default": if_false,
-     })
- 
\ No newline at end of file
diff --git a/third_party/xla_client/BUILD b/third_party/xla_client/BUILD
index 77a8ce29b9e5..996d0b901a10 100644
--- a/third_party/xla_client/BUILD
+++ b/third_party/xla_client/BUILD
@@ -1,6 +1,5 @@
 load(
     "@org_tensorflow//tensorflow:tensorflow.bzl",
-    "if_with_tpu_support",
     "tf_cc_shared_object",
 )
 load(
@@ -11,6 +10,10 @@ load(
     "@org_tensorflow//tensorflow/tsl/platform/default:cuda_build_defs.bzl",
     "if_cuda_is_configured",
 )
+load(
+    "//bazel:tensorflow.bzl",
+    "if_with_tpu_support",
+)
 
 licenses(["notice"])  # Apache 2.0