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

Fix build breakage from OpenXLA switch.

Browse Source 
PiperOrigin-RevId: 516325478
This commit is contained in:
Peter Hawkins 2023-03-13 14:35:56 -07:00 committed by jax authors
parent 5aa74acbcd
commit ab45383038
4 changed files with 40 additions and 40 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
examples/jax_cpp/BUILD
View File

@ -19,15 +19,15 @@ cc_binary(
srcs = ["main.cc"],
tags = ["manual"],
deps = [
"@xla//:literal",
"@xla//:literal_util",
"@xla//:shape_util",
"@xla//:status",
"@xla//:statusor",
"@xla///pjrt:pjrt_client",
"@xla///pjrt:tfrt_cpu_pjrt_client",
"@xla///service:hlo_proto_cc",
"@xla///tools:hlo_module_loader",
"@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:tfrt_cpu_pjrt_client",
"@xla//xla/service:hlo_proto_cc",
"@xla//xla/tools:hlo_module_loader",
"@tsl///platform:logging",
"@tsl///platform:platform_port",
],

6
jaxlib/BUILD
View File

@ -68,7 +68,7 @@ symlink_files(
symlink_files(
name = "xla_client",
srcs = ["@xla///python:xla_client"],
srcs = ["@xla//xla/python:xla_client"],
dst = ".",
flatten = True,
)
@ -76,8 +76,8 @@ symlink_files(
symlink_files(
name = "xla_extension",
srcs = if_windows(
["@xla///python:xla_extension.pyd"],
["@xla///python:xla_extension.so"],
["@xla//xla/python:xla_extension.pyd"],
["@xla//xla/python:xla_extension.so"],
),
dst = ".",
flatten = True,

6
jaxlib/cpu/BUILD
View File

@ -31,7 +31,7 @@ cc_library(
srcs = ["lapack_kernels.cc"],
hdrs = ["lapack_kernels.h"],
deps = [
"@xla///service:custom_call_status",
"@xla//xla/service:custom_call_status",
"@com_google_absl//absl/base:dynamic_annotations",
],
)
@ -74,7 +74,7 @@ cc_library(
features = ["-use_header_modules"],
deps = [
":ducc_fft_flatbuffers_cc",
"@xla///service:custom_call_status",
"@xla//xla/service:custom_call_status",
"@ducc",
"@flatbuffers//:runtime_cc",
],
@ -106,7 +106,7 @@ cc_library(
":ducc_fft_kernels",
":lapack_kernels",
":lapack_kernels_using_lapack",
"@xla///service:custom_call_target_registry",
"@xla//xla/service:custom_call_target_registry",
],
alwayslink = 1,
)

50
jaxlib/cuda/BUILD
View File

@ -51,8 +51,8 @@ cc_library(
features = ["-use_header_modules"],
deps = [
":cuda_vendor",
"@xla///stream_executor/cuda:cusolver_lib",
"@xla///stream_executor/cuda:cusparse_lib",
"@xla//xla/stream_executor/cuda:cusolver_lib",
"@xla//xla/stream_executor/cuda:cusparse_lib",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
@ -72,9 +72,9 @@ cc_library(
":cuda_vendor",
"//jaxlib:handle_pool",
"//jaxlib:kernel_helpers",
"@xla///service:custom_call_status",
"@xla///stream_executor/cuda:cublas_lib",
"@xla///stream_executor/cuda:cudart_stub",
"@xla//xla/service:custom_call_status",
"@xla//xla/stream_executor/cuda:cublas_lib",
"@xla//xla/stream_executor/cuda:cudart_stub",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/base",
"@com_google_absl//absl/base:core_headers",
@ -102,7 +102,7 @@ pybind_extension(
":cublas_kernels",
":cuda_vendor",
"//jaxlib:kernel_pybind11_helpers",
"@xla///stream_executor/cuda:cublas_lib",
"@xla//xla/stream_executor/cuda:cublas_lib",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/strings:str_format",
"@pybind11",
@ -118,9 +118,9 @@ cc_library(
":cuda_vendor",
"//jaxlib:handle_pool",
"//jaxlib:kernel_helpers",
"@xla///service:custom_call_status",
"@xla///stream_executor/cuda:cudart_stub",
"@xla///stream_executor/cuda:cudnn_lib",
"@xla//xla/service:custom_call_status",
"@xla//xla/stream_executor/cuda:cudart_stub",
"@xla//xla/stream_executor/cuda:cudnn_lib",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings:str_format",
@ -157,8 +157,8 @@ cc_library(
":cuda_vendor",
"//jaxlib:handle_pool",
"//jaxlib:kernel_helpers",
"@xla///service:custom_call_status",
"@xla///stream_executor/cuda:cusolver_lib",
"@xla//xla/service:custom_call_status",
"@xla//xla/stream_executor/cuda:cusolver_lib",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/synchronization",
@ -180,8 +180,8 @@ pybind_extension(
":cuda_vendor",
":cusolver_kernels",
"//jaxlib:kernel_pybind11_helpers",
"@xla///stream_executor/cuda:cudart_stub",
"@xla///stream_executor/cuda:cusolver_lib",
"@xla//xla/stream_executor/cuda:cudart_stub",
"@xla//xla/stream_executor/cuda:cusolver_lib",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/strings:str_format",
"@local_config_cuda//cuda:cuda_headers",
@ -198,9 +198,9 @@ cc_library(
":cuda_vendor",
"//jaxlib:handle_pool",
"//jaxlib:kernel_helpers",
"@xla///service:custom_call_status",
"@xla///stream_executor/cuda:cudart_stub",
"@xla///stream_executor/cuda:cusparse_lib",
"@xla//xla/service:custom_call_status",
"@xla//xla/stream_executor/cuda:cudart_stub",
"@xla//xla/stream_executor/cuda:cusparse_lib",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/synchronization",
@ -222,8 +222,8 @@ pybind_extension(
":cuda_vendor",
":cusparse_kernels",
"//jaxlib:kernel_pybind11_helpers",
"@xla///stream_executor/cuda:cudart_stub",
"@xla///stream_executor/cuda:cusparse_lib",
"@xla//xla/stream_executor/cuda:cudart_stub",
"@xla//xla/stream_executor/cuda:cusparse_lib",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/base",
"@com_google_absl//absl/base:core_headers",
@ -249,7 +249,7 @@ cc_library(
":cuda_lu_pivot_kernels_impl",
":cuda_vendor",
"//jaxlib:kernel_helpers",
"@xla///service:custom_call_status",
"@xla//xla/service:custom_call_status",
"@local_config_cuda//cuda:cuda_headers",
],
)
@ -264,7 +264,7 @@ cuda_library(
":cuda_gpu_kernel_helpers",
":cuda_vendor",
"//jaxlib:kernel_helpers",
"@xla///service:custom_call_status",
"@xla//xla/service:custom_call_status",
"@local_config_cuda//cuda:cuda_headers",
],
)
@ -284,7 +284,7 @@ pybind_extension(
":cuda_lu_pivot_kernels_impl",
":cuda_vendor",
"//jaxlib:kernel_pybind11_helpers",
"@xla///stream_executor/cuda:cudart_stub",
"@xla//xla/stream_executor/cuda:cudart_stub",
"@local_config_cuda//cuda:cuda_headers",
"@pybind11",
],
@ -301,7 +301,7 @@ cc_library(
":cuda_prng_kernels_impl",
":cuda_vendor",
"//jaxlib:kernel_helpers",
"@xla///service:custom_call_status",
"@xla//xla/service:custom_call_status",
"@local_config_cuda//cuda:cuda_headers",
],
)
@ -316,7 +316,7 @@ cuda_library(
":cuda_gpu_kernel_helpers",
":cuda_vendor",
"//jaxlib:kernel_helpers",
"@xla///service:custom_call_status",
"@xla//xla/service:custom_call_status",
"@local_config_cuda//cuda:cuda_headers",
],
)
@ -334,7 +334,7 @@ pybind_extension(
":cuda_gpu_kernel_helpers",
":cuda_prng_kernels",
"//jaxlib:kernel_pybind11_helpers",
"@xla///stream_executor/cuda:cudart_stub",
"@xla//xla/stream_executor/cuda:cudart_stub",
"@local_config_cuda//cuda:cuda_headers",
"@pybind11",
],
@ -351,7 +351,7 @@ cc_library(
":cuda_vendor",
":cusolver_kernels",
":cusparse_kernels",
"@xla///service:custom_call_target_registry",
"@xla//xla/service:custom_call_target_registry",
],
alwayslink = 1,
)