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

Add --config=cuda_libraries_from_stubs in the end of all additional bazel options for CUDA wheels.

Browse Source 
Build.py shouldn't be used for building the wheels with real CUDA libraries in the dependencies. This change prevents overriding the default configuration.

PiperOrigin-RevId: 725326252
This commit is contained in:
jax authors 2025-02-10 13:24:09 -08:00
parent b7d012281e
commit e64650e2ba
1 changed files with 2 additions and 0 deletions

2
build/build.py

@ -588,6 +588,8 @@ async def main():
)
for option in args.bazel_options:
wheel_build_command_base.append(option)
if "cuda" in args.wheels:
wheel_build_command_base.append("--config=cuda_libraries_from_stubs")
with open(".jax_configure.bazelrc", "w") as f:
jax_configure_options = utils.get_jax_configure_bazel_options(wheel_build_command_base.get_command_as_list())