Move stable hlo compile test to XLA:CPU public API

PiperOrigin-RevId: 698121993
openxla · Nov 19, 2024 · 1edbb92 · 1edbb92
1 parent 36c6daf
commit 1edbb92
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/xla/examples/axpy/BUILD b/xla/examples/axpy/BUILD
@@ -16,7 +16,8 @@ xla_cc_test(
         "//xla/pjrt:pjrt_executable",
         "//xla/pjrt/c:pjrt_c_api_cpu",
         "//xla/pjrt/c:pjrt_c_api_hdrs",
-        "//xla/pjrt/cpu:cpu_client",
+        "//xla/pjrt/plugin/xla_cpu:cpu_client_options",
+        "//xla/pjrt/plugin/xla_cpu:xla_cpu_pjrt_client",
         "//xla/tests:literal_test_util",
         "//xla/tests:xla_internal_test_main",
         "//xla/tsl/lib/core:status_test_util",

diff --git a/xla/examples/axpy/stablehlo_compile_test.cc b/xla/examples/axpy/stablehlo_compile_test.cc
@@ -35,11 +35,12 @@ limitations under the License.
 #include "xla/literal_util.h"
 #include "xla/pjrt/c/pjrt_c_api.h"
 #include "xla/pjrt/c/pjrt_c_api_cpu.h"
-#include "xla/pjrt/cpu/cpu_client.h"
 #include "xla/pjrt/pjrt_api.h"
 #include "xla/pjrt/pjrt_c_api_client.h"
 #include "xla/pjrt/pjrt_client.h"
 #include "xla/pjrt/pjrt_executable.h"
+#include "xla/pjrt/plugin/xla_cpu/cpu_client_options.h"
+#include "xla/pjrt/plugin/xla_cpu/xla_cpu_pjrt_client.h"
 #include "xla/tests/literal_test_util.h"
 #include "xla/tsl/lib/core/status_test_util.h"
 #include "tsl/platform/env.h"
@@ -129,7 +130,7 @@ TEST_F(StableHloAxpyTest, CompileAndExecuteCPUTestProgram) {
   xla::CpuClientOptions options;
   options.cpu_device_count = 4;
   TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<xla::PjRtClient> client,
-                          xla::GetTfrtCpuClient(std::move(options)));
+                          xla::GetXlaPjrtCpuClient(options));
 
   TF_ASSERT_OK_AND_ASSIGN(mlir::OwningOpRef<mlir::ModuleOp> program,
                           CreateStableHloProgram(GetTestProgramPath()));