Package Details: python-jaxlib-cuda 0.4.38-1

Git Clone URL: https://aur.archlinux.org/python-jaxlib-cuda.git (read-only, click to copy)
Package Base: python-jaxlib-cuda
Description: XLA library for JAX
Upstream URL: https://github.com/jax-ml/jax/
Keywords: deep-learning google jax maching-learning xla
Licenses: Apache
Groups: jax
Conflicts: python-jaxlib
Provides: python-jaxlib
Submitter: daskol
Maintainer: daskol
Last Packager: daskol
Votes: 9
Popularity: 0.40
First Submitted: 2023-02-12 23:18 (UTC)
Last Updated: 2024-12-24 19:26 (UTC)

Latest Comments

1 2 3 4 5 Next › Last »

truncs commented on 2025-03-05 17:07 (UTC)

@medaminezghal This works! Thanks for sharing the PKGBUILD file

medaminezghal commented on 2025-03-01 16:14 (UTC)

@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
}

medaminezghal commented on 2025-02-27 16:52 (UTC)

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

medaminezghal commented on 2025-02-12 13:02 (UTC) (edited on 2025-02-12 13:59 (UTC) by medaminezghal)

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

medaminezghal commented on 2025-02-10 07:39 (UTC)

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

eterhskip commented on 2025-01-23 18:20 (UTC)

Package is not building for me.

actionless commented on 2025-01-02 17:06 (UTC)

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' \

daskol commented on 2024-12-26 01:14 (UTC) (edited on 2024-12-26 01:37 (UTC) by daskol)

@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.

medaminezghal commented on 2024-12-25 14:13 (UTC) (edited on 2024-12-25 15:08 (UTC) by medaminezghal) 

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?

actionless commented on 2024-11-10 13:48 (UTC)

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

1 2 3 4 5 Next › Last »