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

Add a recently added Mosaic Python file to build_wheel.py

Browse Source 
PiperOrigin-RevId: 584356541
This commit is contained in:
Adam Paszke 2023-11-21 10:03:15 -08:00 committed by jax authors
parent 9c58f30bcc
commit 038879248d
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

@ -249,6 +249,7 @@ def prepare_wheel(sources_path: pathlib.Path, *, cpu, include_gpu_plugin_extensi
src_files=[
"__main__/jaxlib/mosaic/python/apply_vector_layout.py",
"__main__/jaxlib/mosaic/python/infer_memref_layout.py",
"__main__/jaxlib/mosaic/python/layout_defs.py",
"__main__/jaxlib/mosaic/python/tpu.py",
],
)