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

Automated Code Change

Browse Source 
PiperOrigin-RevId: 635818645
This commit is contained in:
Kyle Lucke 2024-05-21 08:39:41 -07:00 committed by jax authors
parent 4394bdc2ad
commit 418b68828a
4 changed files with 2 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
examples/jax_cpp/BUILD
View File

@ -23,12 +23,9 @@ cc_binary(
deps = [
"@xla//xla:literal",
"@xla//xla:literal_util",
"@xla//xla:shape_util",
"@xla//xla:status",
"@xla//xla:statusor",
"@xla//xla/pjrt:pjrt_client",
"@xla//xla/pjrt/cpu:cpu_client",
"@xla//xla/service:hlo_proto_cc",
"@xla//xla/tools:hlo_module_loader",
"@tsl//tsl/platform:logging",
"@tsl//tsl/platform:platform_port",

1
examples/jax_cpp/main.cc
View File

@ -44,7 +44,6 @@ limitations under the License.
#include "xla/literal_util.h"
#include "xla/pjrt/cpu/cpu_client.h"
#include "xla/pjrt/pjrt_client.h"
#include "xla/status.h"
#include "xla/statusor.h"
#include "xla/tools/hlo_module_loader.h"
#include "tsl/platform/init_main.h"

1
jaxlib/BUILD
View File

@ -187,6 +187,7 @@ pybind_extension(
srcs = ["cuda_plugin_extension.cc"],
module_name = "cuda_plugin_extension",
deps = [
"@com_google_absl//absl/status",
"@nanobind",
"//jaxlib:kernel_nanobind_helpers",
"@xla//third_party/python_runtime:headers",

2
jaxlib/cuda_plugin_extension.cc
View File

@ -19,6 +19,7 @@ limitations under the License.
#include <utility>
#include "nanobind/nanobind.h"
#include "absl/status/status.h"
#include "third_party/gpus/cuda/include/cuda.h"
#include "jaxlib/kernel_nanobind_helpers.h"
#include "xla/pjrt/c/pjrt_c_api.h"
@ -26,7 +27,6 @@ limitations under the License.
#include "xla/pjrt/c/pjrt_c_api_helpers.h"
#include "xla/pjrt/status_casters.h"
#include "xla/python/py_client_gpu.h"
#include "xla/status.h"
#include "xla/tsl/python/lib/core/numpy.h"
#include "xla/util.h"