Building from Source

[raisin]

When following the build from source guide at the url:

https://jax.readthedocs.io/en/latest/developer.html

I encounter an issue with the command:

python build/build.py build --wheels=jaxlib --verbose

The flag "--wheels=jaxlib" doesn't seem to be a valid argument for build.py. The shell output is pasted below. The command is run in a Conda environment with CUDA 12.5 and CUDNN 9.3 on the jax repo with tag jax-v0.4.32. Furthermore, the Python version is 3.12.8.

(JAX-0432-CUDA-1251) ➜ jax git:(jax-v0.4.32) ✗ python build/build.py build --wheels=jaxlib --verbose

usage: build.py [-h] [--verbose [VERBOSE] | --noverbose] [--bazel_path BAZEL_PATH]
[--python_bin_path PYTHON_BIN_PATH] [--target_cpu_features {release,native,default}]
[--use_clang [USE_CLANG] | --nouse_clang] [--clang_path CLANG_PATH]
[--enable_mkl_dnn [ENABLE_MKL_DNN] | --noenable_mkl_dnn]
[--enable_cuda [ENABLE_CUDA] | --noenable_cuda]
[--build_gpu_plugin [BUILD_GPU_PLUGIN] | --nobuild_gpu_plugin]
[--build_gpu_kernel_plugin {cuda,rocm}]
[--build_gpu_pjrt_plugin [BUILD_GPU_PJRT_PLUGIN] | --nobuild_gpu_pjrt_plugin]
[--gpu_plugin_cuda_version {12}] [--gpu_plugin_rocm_version {60}]
[--enable_rocm [ENABLE_ROCM] | --noenable_rocm] [--enable_nccl [ENABLE_NCCL] |
--noenable_nccl] [--remote_build [REMOTE_BUILD] | --noremote_build]
[--cuda_version CUDA_VERSION] [--cudnn_version CUDNN_VERSION]
[--cuda_compute_capabilities CUDA_COMPUTE_CAPABILITIES]
[--rocm_amdgpu_targets ROCM_AMDGPU_TARGETS] [--rocm_path ROCM_PATH]
[--bazel_startup_options BAZEL_STARTUP_OPTIONS] [--bazel_options BAZEL_OPTIONS]
[--output_path OUTPUT_PATH] [--target_cpu TARGET_CPU] [--editable]
[--python_version PYTHON_VERSION] [--configure_only [CONFIGURE_ONLY] |
--noconfigure_only] [--requirements_update [REQUIREMENTS_UPDATE] |
--norequirements_update]
[--requirements_nightly_update [REQUIREMENTS_NIGHTLY_UPDATE] |
--norequirements_nightly_update]
build.py: error: unrecognized arguments: build --wheels=jaxlib

NOTE: I searched for similar discussions. Based on those discussions, it seems people were building with alternative flags. However, it would be good if the docs were updated.

[raisin]

From a discussion perspective, it would be good to know:

1. Are the docs at readthedocs upto date and generally reliable for JAX? If yes, is the issue above an error in the docs?
2. If readthedocs is not upto date, is there an alternative documentation source that people generally use, which is better?

[dfm]

Thanks for the question! Yes, I believe the readthedocs page is up to date, but those docs are for the current version of JAX (0.4.38), and I wouldn't necessarily expect those instructions to work with older versions like the one you're trying (0.4.32). Unfortunately we don't have an archive of the rendered docs from that version, but you could take a look at the docs/developer.md file in the JAX repo at that tag for the state when that release was made. Hope that helps!

[raisin]

Thanks - It's helpful to know those instructions work with 0.4.38. After switching to that tag, I was able to start the build process.