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

Make //jax:tpu_custom_call respect --//jax:build_jaxlib=false

Browse Source 
Otherwise jaxlib is partially built and doesn't work properly.
This commit is contained in:
Skye Wanderman-Milne 2023-07-26 15:50:42 +00:00
parent f66d3cf016
commit d0b65f2ab8
1 changed files with 9 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
jax/BUILD
View File

@ -572,11 +572,15 @@ pytype_strict_library(
":core",
":jax",
"//jax/_src/lib",
"//jaxlib/mlir:ir",
"//jaxlib/mlir:mhlo_dialect",
"//jaxlib/mlir:pass_manager",
"//jaxlib/mlir:stablehlo_dialect",
] + py_deps("numpy") + py_deps("absl/flags"),
] + select({
":enable_jaxlib_build": [
"//jaxlib/mlir:ir",
"//jaxlib/mlir:mhlo_dialect",
"//jaxlib/mlir:pass_manager",
"//jaxlib/mlir:stablehlo_dialect",
],
"//conditions:default": [],
}) + py_deps("numpy") + py_deps("absl/flags"),
)
pytype_strict_library(