From ee541f2cf6c525d326cfb6bfe4d30ee38361c56d Mon Sep 17 00:00:00 2001
From: Qiao Zhang <zhangqiaorjc@google.com>
Date: Fri, 15 Jan 2021 18:40:02 -0800
Subject: [PATCH] Internal changes only.

PiperOrigin-RevId: 352121854
Change-Id: I9e7679109d15dd4a8ec75cfc295c18dd3c1eedff
---
 tensorflow/compiler/xla/pjrt/BUILD | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/tensorflow/compiler/xla/pjrt/BUILD b/tensorflow/compiler/xla/pjrt/BUILD
index b368afd88e6..145f5f71092 100644
--- a/tensorflow/compiler/xla/pjrt/BUILD
+++ b/tensorflow/compiler/xla/pjrt/BUILD
@@ -8,6 +8,16 @@ package(
     licenses = ["notice"],  # Apache 2.0
 )
 
+package_group(
+    name = "friends",
+    includes = [
+        "//tensorflow:internal",
+    ],
+    packages = [
+        "//third_party/py/jax/...",
+    ],
+)
+
 cc_library(
     name = "worker_thread",
     srcs = ["worker_thread.cc"],
@@ -274,6 +284,7 @@ cc_library(
     name = "cpu_device",
     srcs = ["cpu_device.cc"],
     hdrs = ["cpu_device.h"],
+    visibility = [":friends"],
     deps = [
         ":pjrt_stream_executor_client",
         "//tensorflow/compiler/xla:statusor",