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

[PJRT C API] Add "cuda_plugin_extension" to "gpu_only_test_deps" to support bazel test for GPU plugin.

Browse Source 
PiperOrigin-RevId: 573251982
This commit is contained in:
Jieying Luo 2023-10-13 10:11:37 -07:00 committed by jax authors
parent 1ae2bbc253
commit 7478fbcfd5
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
jaxlib/cuda/BUILD
View File

@ -514,5 +514,8 @@ py_library(
py_library(
name = "gpu_only_test_deps",
# `if_cuda_is_configured` will default to `[]`.
deps = if_cuda_is_configured([":cuda_gpu_support"]),
deps = if_cuda_is_configured([
":cuda_gpu_support",
"//jaxlib:cuda_plugin_extension",
]),
)