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

Use tensorflow/compiler/xla/stream_executor instead of tensorflow/stream_executor

Browse Source 
PiperOrigin-RevId: 470804752
This commit is contained in:
David Dunleavy 2022-08-29 13:45:45 -07:00 committed by jax authors
parent bff98655a1
commit a8aa774a57
1 changed files with 14 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

28
jaxlib/cuda/BUILD
View File

@ -33,8 +33,8 @@ cc_library(
],
features = ["-use_header_modules"],
deps = [
"@org_tensorflow//tensorflow/stream_executor/cuda:cusolver_lib",
"@org_tensorflow//tensorflow/stream_executor/cuda:cusparse_lib",
"@org_tensorflow//tensorflow/compiler/xla/stream_executor/cuda:cusolver_lib",
"@org_tensorflow//tensorflow/compiler/xla/stream_executor/cuda:cusparse_lib",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
@ -54,8 +54,8 @@ cc_library(
"//jaxlib:handle_pool",
"//jaxlib:kernel_helpers",
"@org_tensorflow//tensorflow/compiler/xla/service:custom_call_status",
"@org_tensorflow//tensorflow/stream_executor/cuda:cublas_lib",
"@org_tensorflow//tensorflow/stream_executor/cuda:cudart_stub",
"@org_tensorflow//tensorflow/compiler/xla/stream_executor/cuda:cublas_lib",
"@org_tensorflow//tensorflow/compiler/xla/stream_executor/cuda:cudart_stub",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/base",
"@com_google_absl//absl/base:core_headers",
@ -82,7 +82,7 @@ pybind_extension(
deps = [
":cublas_kernels",
"//jaxlib:kernel_pybind11_helpers",
"@org_tensorflow//tensorflow/stream_executor/cuda:cublas_lib",
"@org_tensorflow//tensorflow/compiler/xla/stream_executor/cuda:cublas_lib",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/strings:str_format",
"@local_config_cuda//cuda:cuda_headers",
@ -99,7 +99,7 @@ cc_library(
"//jaxlib:handle_pool",
"//jaxlib:kernel_helpers",
"@org_tensorflow//tensorflow/compiler/xla/service:custom_call_status",
"@org_tensorflow//tensorflow/stream_executor/cuda:cusolver_lib",
"@org_tensorflow//tensorflow/compiler/xla/stream_executor/cuda:cusolver_lib",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/synchronization",
@ -120,8 +120,8 @@ pybind_extension(
":cuda_gpu_kernel_helpers",
":cusolver_kernels",
"//jaxlib:kernel_pybind11_helpers",
"@org_tensorflow//tensorflow/stream_executor/cuda:cudart_stub",
"@org_tensorflow//tensorflow/stream_executor/cuda:cusolver_lib",
"@org_tensorflow//tensorflow/compiler/xla/stream_executor/cuda:cudart_stub",
"@org_tensorflow//tensorflow/compiler/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",
@ -138,8 +138,8 @@ cc_library(
"//jaxlib:handle_pool",
"//jaxlib:kernel_helpers",
"@org_tensorflow//tensorflow/compiler/xla/service:custom_call_status",
"@org_tensorflow//tensorflow/stream_executor/cuda:cudart_stub",
"@org_tensorflow//tensorflow/stream_executor/cuda:cusparse_lib",
"@org_tensorflow//tensorflow/compiler/xla/stream_executor/cuda:cudart_stub",
"@org_tensorflow//tensorflow/compiler/xla/stream_executor/cuda:cusparse_lib",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/synchronization",
@ -160,8 +160,8 @@ pybind_extension(
":cuda_gpu_kernel_helpers",
":cusparse_kernels",
"//jaxlib:kernel_pybind11_helpers",
"@org_tensorflow//tensorflow/stream_executor/cuda:cudart_stub",
"@org_tensorflow//tensorflow/stream_executor/cuda:cusparse_lib",
"@org_tensorflow//tensorflow/compiler/xla/stream_executor/cuda:cudart_stub",
"@org_tensorflow//tensorflow/compiler/xla/stream_executor/cuda:cusparse_lib",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/base",
"@com_google_absl//absl/base:core_headers",
@ -219,7 +219,7 @@ pybind_extension(
":cuda_lu_pivot_kernels",
":cuda_lu_pivot_kernels_impl",
"//jaxlib:kernel_pybind11_helpers",
"@org_tensorflow//tensorflow/stream_executor/cuda:cudart_stub",
"@org_tensorflow//tensorflow/compiler/xla/stream_executor/cuda:cudart_stub",
"@local_config_cuda//cuda:cuda_headers",
"@pybind11",
],
@ -267,7 +267,7 @@ pybind_extension(
":cuda_gpu_kernel_helpers",
":cuda_prng_kernels",
"//jaxlib:kernel_pybind11_helpers",
"@org_tensorflow//tensorflow/stream_executor/cuda:cudart_stub",
"@org_tensorflow//tensorflow/compiler/xla/stream_executor/cuda:cudart_stub",
"@local_config_cuda//cuda:cuda_headers",
"@pybind11",
],