The set of editable wheels (`jaxlib`, `jax-cuda-plugin` and `jax-cuda-pjrt`) was used as dependencies in `requirements.in` file together with `:build_jaxlib=false` flag.
After [adding `jax` wheel dependencies](f5a4d1a85c) to the tests when `:build_jaxlib=false` is used, we need an editable `jax` wheel target as well to get the tests passing.
PiperOrigin-RevId: 740840736
Remove jax dependencies from the Bazel test targets for `:build_jaxlib=false` and `:build_jaxlib=wheel`.
`internal_test_util` is removed from the `jax` wheel. To use this package in Bazel py_test, we need to copy it to the unpacked wheel folder. This is done by adding `wheel_deps` value to `py_import` Jax targets.
This change concludes ML Wheels design implementation in JAX and enables testing of all wheels via Bazel command.
PiperOrigin-RevId: 740037952
Refactor `jax_wheel` macros, so it outputs a `.whl` file only.
When the macros returns one output object only, it allows all downstream dependencies consume it easily without the need to filter the macros outputs.
The previous implementation design (when `jax_wheel` returned `.tar.gz` and `.whl` files) required one of two options: either create a new target that produces `.whl` only, or to implement filename filtering in the downstream rules. With the new implementation we can just depend on `//:jax_wheel` target that produces the `.whl`.
PiperOrigin-RevId: 738547491
This change replicates the old method of building `jax` wheel via `python -m build`, which produced `.tar.gz` and `.whl` files.
PiperOrigin-RevId: 731721522
Change in preparation for supporting running JAX tests under Bazel. This change allows the Bazel py_library() to see version.py.
Update symlink_files Bazel macro to a newer version.
PiperOrigin-RevId: 458481396
Currently jax and jaxlib have separate version numbers in the JAX source
tree. It is tedious and confusing to bump both version numbers.
However, there is a simpler way to think of things: it is the source
tree that is versioned using a single version number, and jax/jaxlib
releases are made using that unified source version number.
PiperOrigin-RevId: 458041752
