jax

6b096b0c - Use common set of build options when building jaxlib+plugin artifacts together

Use common set of build options when building jaxlib+plugin artifacts together This commit modifies the behavior of the build CLI when building jaxlib and GPU plugin artifacts together (for instance `python build --wheels=jaxlib,jax-cuda-plugin`. Before, CUDA/ROCm build options were only passed when building the CUDA/ROCm artifacts. However, this leads to inefficient use of the build cache as it looks like Bazel tries to rebuild some targets that has already been built in the previous run. This seems to be because the GPU plugin artifacts have a different set of build options compared to `jaxlib` which for some reason causes Bazel to invalidate/ignore certain cache hits. Therefore, this commit makes it so that the build options remain the same when the `jaxlib` and GPU artifacts are being built together so that we can better utilize the build cache. As an example, this means that if `python build --wheels=jaxlib,jax-cuda-plugin` is run, the following build options will apply to both `jaxlib` and `jax-cuda-plugin` builds: ``` /usr/local/bin/bazel run --repo_env=HERMETIC_PYTHON_VERSION=3.10 \ --verbose_failures=true --action_env=CLANG_COMPILER_PATH="/usr/lib/llvm-16/bin/clang" \ --repo_env=CC="/usr/lib/llvm-16/bin/clang" \ --repo_env=BAZEL_COMPILER="/usr/lib/llvm-16/bin/clang" \ --config=clang --config=mkl_open_source_only --config=avx_posix \ --config=cuda --action_env=CLANG_CUDA_COMPILER_PATH="/usr/lib/llvm-16/bin/clang" \ --config=build_cuda_with_nvcc ``` Note, this commit shouldn't affect the content of the wheel it self. It is only meant to give a performance boost when building `jalxib`+plugin aritfacts together. Also, this removes code that was used to build (now deprecated) monolithic `jaxlib` build from `build_wheel.py` PiperOrigin-RevId: 708035062