Move some tests to public XLA:CPU API

PiperOrigin-RevId: 698116537
openxla · Nov 19, 2024 · e191661 · e191661
1 parent 7fd2196
commit e191661
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/xla/pjrt/BUILD b/xla/pjrt/BUILD
@@ -697,15 +697,15 @@ cc_library(
     ],
 )
 
-# Transitional forwarding target. Use cpu:cpu_client instead.
+# Transitional forwarding target. Use pjrt/plugin/xla_cpu:xla_cpu_pjrt_client instead.
 cc_library(
     name = "tfrt_cpu_pjrt_client",
     hdrs = ["tfrt_cpu_pjrt_client.h"],
     visibility = internal_visibility([
         "//xla:friends",
     ]),
     deps = [
-        "//xla/pjrt/cpu:cpu_client",
+        "//xla/pjrt/plugin/xla_cpu:xla_cpu_pjrt_client",
     ],
 )
 

diff --git a/xla/pjrt/tfrt_cpu_pjrt_client.h b/xla/pjrt/tfrt_cpu_pjrt_client.h
@@ -18,6 +18,6 @@ limitations under the License.
 
 // Transitional forwarding header. Please include cpu/cpu_client.h directly.
 
-#include "xla/pjrt/cpu/cpu_client.h"
+#include "xla/pjrt/plugin/xla_cpu/xla_cpu_pjrt_client.h"
 
 #endif  // XLA_PJRT_TFRT_CPU_PJRT_CLIENT_H_