rocm_jax

mirror of https://github.com/ROCm/jax.git synced 2025-04-16 11:56:07 +00:00

History

Jieying Luo 9c2caedab1 Add subdirectories to the output path when building editable wheels for jaxlib and GPU plugin.

When `build_gpu_plugin` is true, three wheels will be produced (jaxlib, jax-cuda-pjrt and jax-cuda-plugin). If they are editable, they need to be placed in subdirectories to avoid overwrite.

Tested on GPU. After the editable wheels are built, they can be installed with `pip install -e /jax/dist/jax_gpu_pjrt /jax/dist/jaxlib /jax/dist/jax_gpu_plugin`.

PiperOrigin-RevId: 660984311

2024-08-08 14:34:55 -07:00

rocm

Fix lint in run_single_gpu.py

2024-08-02 16:03:03 -07:00

BUILD.bazel

Use hermetic Python in JAX, see "Managing hermetic Python" in developer.md for details

2024-05-15 18:20:56 -07:00

build.py

Add subdirectories to the output path when building editable wheels for jaxlib and GPU plugin.

2024-08-08 14:34:55 -07:00

collect-profile-requirements.txt

Cloud TPU CI: make sure we update test deps and upgrade protobuf version