Internal changes only.

PiperOrigin-RevId: 352121854
Change-Id: I9e7679109d15dd4a8ec75cfc295c18dd3c1eedff

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",