Arch Linux User Repository

@medaminezghal This works! Thanks for sharing the PKGBUILD file

@daskol I have successfully built the package using this PKGBUILD file:

# Maintainer: Daniel Bershatsky <bepshatsky@yandex.ru>

pkgname=python-jaxlib-cuda
pkgver=0.5.1
pkgrel=1
pkgdesc='XLA library for JAX'
arch=('x86_64')
url='https://github.com/jax-ml/jax/'
license=('Apache')
groups=('jax')
depends=('cuda' 'cudnn' 'nccl' 'openssl' 'python-absl' 'python-flatbuffers'
         'python-ml-dtypes>=0.4.0' 'python-numpy' 'python-scipy')
makedepends=('clang18' 'python-build' 'python-installer' 'python-setuptools' 'python-wheel')
conflicts=('python-jaxlib')
provides=("python-jaxlib=$pkgver")
source=("jaxlib-${pkgver}.tar.gz::$url/archive/refs/tags/jax-v${pkgver}.tar.gz"
        'https://github.com/bazelbuild/bazel/releases/download/6.5.0/bazel-6.5.0-linux-x86_64')
sha256sums=('e74b1209517682075933f757d646b73040d09fe39ee3e9e4cd398407dd0902d2'
            'a40ac69263440761199fcb8da47ad4e3f328cbe79ffbf4ecc14e5ba252857307')

prepare() {
    ln -sf $(readlink bazel-6.5.0-linux-x86_64) $srcdir/jax-jax-v${pkgver}/build
    chmod +x $srcdir/bazel-6.5.0-linux-x86_64
}

build() {
    # CUDA and CUDNN versions.
    export CUDA_VERSION=$(/opt/cuda/bin/nvcc --version | sed -n 's/^.*release \(.*\),.*/\1/p').0
    export CUDNN_MAJOR=$(sed -n 's/^#define CUDNN_MAJOR\s*\(.*\).*/\1/p' /usr/include/cudnn_version.h)
    export CUDNN_MINOR=$(sed -n 's/^#define CUDNN_MINOR\s*\(.*\).*/\1/p' /usr/include/cudnn_version.h)
    export CUDNN_PATCH=$(sed -n 's/^#define CUDNN_PATCHLEVEL\s*\(.*\).*/\1/p' /usr/include/cudnn_version.h)
    export CUDNN_VERSION="${CUDNN_MAJOR}.${CUDNN_MINOR}.${CUDNN_PATCH}"
    export CUDA_COMPUTE_CAPABILITIES=sm_80,sm_86,sm_89,sm_90,compute_90

    # Override default version.
    export JAXLIB_RELEASE=$pkgver

    cd $srcdir/jax-jax-v$pkgver
    build/build.py build \
        --bazel_options='--action_env=JAXLIB_RELEASE' \
        --bazel_options="--action_env=TF_CUDA_COMPUTE_CAPABILITIES=${CUDA_COMPUTE_CAPABILITIES}" \
        --bazel_options="--repo_env=HERMETIC_CUDA_VERSION=${CUDA_VERSION}" \
        --bazel_options='--repo_env=LOCAL_NCCL_PATH=/usr' \
        --bazel_path=$srcdir/bazel-6.5.0-linux-x86_64 \
        --bazel_startup_options="--output_user_root=$srcdir/bazel"\
    --build_cuda_with_clang \
    --use_clang --clang_path='/usr/lib/llvm18/bin/clang' \
        --cuda_version=$CUDA_VERSION \
        --cudnn_version=$CUDNN_VERSION \
    --target_cpu_features=release \
    --verbose \
        --wheels=jaxlib,jax-cuda-plugin,jax-cuda-pjrt
}

package() {
    cd $srcdir/jax-jax-v$pkgver
    python -m installer --compile-bytecode=1 --destdir=$pkgdir \
        $srcdir/jax-jax-v$pkgver/dist/jaxlib-$pkgver-*.whl
    python -m installer --compile-bytecode=1 --destdir=$pkgdir \
        $srcdir/jax-jax-v$pkgver/dist/jax_cuda12_pjrt-$pkgver-*.whl
    python -m installer --compile-bytecode=1 --destdir=$pkgdir \
        $srcdir/jax-jax-v$pkgver/dist/jax_cuda12_plugin-$pkgver-*.whl
}

@daskol could you update this package to the latest version?

@daskol I have successfully built the jaxlib-cuda 0.5.0 using this fix and this file from this commit.

@daskol do we need to wait until the xla be compatible with CUDA 12.8?

Package is not building for me.

it doesn't build if ccache is installed, workaround:

diff --git a/PKGBUILD b/PKGBUILD
index e4cec07..3c972bb 100644
--- a/PKGBUILD
+++ b/PKGBUILD
@@ -35,6 +35,8 @@ build() {
     # Override default version.
     export JAXLIB_RELEASE=$pkgver

+       export PATH=$(echo $PATH | tr ":" "\n" | grep -v ccache | tr "\n" ":")
+
     cd $srcdir/jax-jax-v$pkgver
     build/build.py build \
         --bazel_options='--action_env=JAXLIB_RELEASE' \

@medaminezghal You can add the following to your .bashrc or .zshrc or any other esource file of your shell.

export XLA_FLAGS=--xla_gpu_cuda_data_dir=/opt/cuda

Could find any solution in the installation process?

Building with local CUDA is broken in JAX across multiple version. So it takes some time. Patch to upstream repository is ready.

W external/xla/xla/service/gpu/llvm_gpu_backend/default/nvptx_libdevice_path.cc:40] Can't find libdevice directory ${CUDA_DIR}/nvvm/libdevice. This may result in compilation or runtime failures, if the program we try to run uses routines from libdevice.
Searched for CUDA in the following directories:
  ./cuda_sdk_lib
  python.runfiles/cuda_nvcc
  python/cuda_nvcc

  /usr/local/cuda
  /usr/lib/python3.13/site-packages/jax_plugins/xla_cuda12/../nvidia/cuda_nvcc
  /usr/lib/python3.13/site-packages/jax_plugins/xla_cuda12/../../nvidia/cuda_nvcc
  /usr/lib/python3.13/site-packages/jax_plugins/xla_cuda12/cuda
  .
You can choose the search directory by setting xla_gpu_cuda_data_dir in HloModule's DebugOptions.  For most apps, setting the environment variable XLA_FLAGS=--xla_gpu_cuda_data_dir=/path/to/cuda will work.

This happen when I test a simple example. To use jax properly, I need to execute export XLA_FLAGS=--xla_gpu_cuda_data_dir=/opt/cuda every time I launch terminal. Could find any solution in the installation process?

it doesn't require any extra files, just update the pkgbuild with the command i described in the github issue