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 gpu_common_utils to build_wheel to fix the gpu wheels build

Browse Source 
PiperOrigin-RevId: 564562958
This commit is contained in:
Yash Katariya 2023-09-11 18:40:17 -07:00 committed by jax authors
parent 76a5dc3cac
commit 2a7b8e6278
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
jaxlib/tools/build_wheel.py
View File

@ -198,6 +198,7 @@ def prepare_wheel(sources_path, *, cpu):
copy_to_jaxlib("__main__/jaxlib/gpu_linalg.py")
copy_to_jaxlib("__main__/jaxlib/gpu_rnn.py")
copy_to_jaxlib("__main__/jaxlib/gpu_triton.py")
copy_to_jaxlib("__main__/jaxlib/gpu_common_utils.py")
copy_to_jaxlib("__main__/jaxlib/gpu_solver.py")
copy_to_jaxlib("__main__/jaxlib/gpu_sparse.py")
copy_to_jaxlib("__main__/jaxlib/tpu_mosaic.py")