mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 11:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

[PJRT] Delete the old :cpu_device target that uses StreamExecutor.

Browse Source 
The TFRT CPU client is better in every way and the SE CPU client is unmaintained and has not been used by JAX in many months.

PiperOrigin-RevId: 489246256
This commit is contained in:
Peter Hawkins 2022-11-17 10:28:19 -08:00 committed by jax authors
parent ebee4f4bfd
commit 88379603e0
2 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
examples/jax_cpp/BUILD
View File

@ -29,8 +29,8 @@ tf_cc_binary(
"@org_tensorflow//tensorflow/compiler/xla:shape_util",
"@org_tensorflow//tensorflow/compiler/xla:status",
"@org_tensorflow//tensorflow/compiler/xla:statusor",
"@org_tensorflow//tensorflow/compiler/xla/pjrt:cpu_device",
"@org_tensorflow//tensorflow/compiler/xla/pjrt:pjrt_client",
"@org_tensorflow//tensorflow/compiler/xla/pjrt:tfrt_cpu_pjrt_client",
"@org_tensorflow//tensorflow/compiler/xla/service:hlo_proto_cc",
"@org_tensorflow//tensorflow/compiler/xla/tools:hlo_module_loader",
"@org_tensorflow//tensorflow/core/platform:logging",

4
examples/jax_cpp/main.cc
View File

@ -42,8 +42,8 @@ limitations under the License.
#include "tensorflow/compiler/xla/literal.h"
#include "tensorflow/compiler/xla/literal_util.h"
#include "tensorflow/compiler/xla/pjrt/cpu_device.h"
#include "tensorflow/compiler/xla/pjrt/pjrt_client.h"
#include "tensorflow/compiler/xla/pjrt/tfrt_cpu_pjrt_client.h"
#include "tensorflow/compiler/xla/status.h"
#include "tensorflow/compiler/xla/statusor.h"
#include "tensorflow/compiler/xla/tools/hlo_module_loader.h"
@ -67,7 +67,7 @@ int main(int argc, char** argv) {
// Get a CPU client.
std::unique_ptr<xla::PjRtClient> client =
xla::GetCpuClient(/*asynchronous=*/true).value();
xla::GetTfrtCpuClient(/*asynchronous=*/true).value();
// Compile XlaComputation to PjRtExecutable.
xla::XlaComputation xla_computation(test_module_proto);