# For the moment we only support the CPU and GPU backends of jaxlib. The TPU # backend will require some additional work. Those wheels are located here: # https://storage.googleapis.com/jax-releases/libtpu_releases.html. # See `python3Packages.jax.passthru` for CUDA tests. { absl-py, autoAddDriverRunpath, autoPatchelfHook, buildPythonPackage, config, fetchPypi, fetchurl, flatbuffers, jaxlib-build, lib, ml-dtypes, python, scipy, stdenv, # Options: cudaSupport ? config.cudaSupport, cudaPackages, }: let inherit (cudaPackages) cudaVersion; version = "0.4.28"; inherit (python) pythonVersion; cudaLibPath = lib.makeLibraryPath ( with cudaPackages; [ (lib.getLib cuda_cudart) # libcudart.so (lib.getLib cuda_cupti) # libcupti.so (lib.getLib cudnn) # libcudnn.so (lib.getLib libcufft) # libcufft.so (lib.getLib libcusolver) # libcusolver.so (lib.getLib libcusparse) # libcusparse.so ] ); # As of 2023-06-06, google/jax upstream is no longer publishing CPU-only wheels to their GCS bucket. Instead the # official instructions recommend installing CPU-only versions via PyPI. cpuSrcs = let getSrcFromPypi = { platform, dist, hash, }: fetchPypi { inherit version platform dist hash ; pname = "jaxlib"; format = "wheel"; # See the `disabled` attr comment below. python = dist; abi = dist; }; in { "3.9-x86_64-linux" = getSrcFromPypi { platform = "manylinux2014_x86_64"; dist = "cp39"; hash = "sha256-Slbr8FtKTBeRaZ2HTgcvP4CPCYa0AQsU+1SaackMqdw="; }; "3.9-aarch64-darwin" = getSrcFromPypi { platform = "macosx_11_0_arm64"; dist = "cp39"; hash = "sha256-sBVi7IrXVxm30DiXUkiel+trTctMjBE75JFjTVKCrTw="; }; "3.9-x86_64-darwin" = getSrcFromPypi { platform = "macosx_10_14_x86_64"; dist = "cp39"; hash = "sha256-T5jMg3srbG3P4Kt/+esQkxSSCUYRmqOvn6oTlxj/J4c="; }; "3.10-x86_64-linux" = getSrcFromPypi { platform = "manylinux2014_x86_64"; dist = "cp310"; hash = "sha256-47zcb45g+FVPQVwU2TATTmAuPKM8OOVGJ0/VRfh1dps="; }; "3.10-aarch64-darwin" = getSrcFromPypi { platform = "macosx_11_0_arm64"; dist = "cp310"; hash = "sha256-8Djmi9ENGjVUcisLvjbmpEg4RDenWqnSg/aW8O2fjAk="; }; "3.10-x86_64-darwin" = getSrcFromPypi { platform = "macosx_10_14_x86_64"; dist = "cp310"; hash = "sha256-pCHSN/jCXShQFm0zRgPGc925tsJvUrxJZwS4eCKXvWY="; }; "3.11-x86_64-linux" = getSrcFromPypi { platform = "manylinux2014_x86_64"; dist = "cp311"; hash = "sha256-Rc4PPIQM/4I2z/JsN/Jsn/B4aV+T4MFiwyDCgfUEEnU="; }; "3.11-aarch64-darwin" = getSrcFromPypi { platform = "macosx_11_0_arm64"; dist = "cp311"; hash = "sha256-eThX+vN/Nxyv51L+pfyBH0NeQ7j7S1AgWERKf17M+Ck="; }; "3.11-x86_64-darwin" = getSrcFromPypi { platform = "macosx_10_14_x86_64"; dist = "cp311"; hash = "sha256-L/gpDtx7ksfq5SUX9lSSYz4mey6QZ7rT5MMj0hPnfPU="; }; "3.12-x86_64-linux" = getSrcFromPypi { platform = "manylinux2014_x86_64"; dist = "cp312"; hash = "sha256-RqGqhX9P7uikP8upXA4Kti1AwmzJcwtsaWVZCLo1n40="; }; "3.12-aarch64-darwin" = getSrcFromPypi { platform = "macosx_11_0_arm64"; dist = "cp312"; hash = "sha256-jdi//jhTcC9jzZJNoO4lc0pNGc1ckmvgM9dyun0cF10="; }; "3.12-x86_64-darwin" = getSrcFromPypi { platform = "macosx_10_14_x86_64"; dist = "cp312"; hash = "sha256-1sCaVFMpciRhrwVuc1FG0sjHTCKsdCaoRetp8ya096A="; }; }; # Note that the prebuilt jaxlib binary requires specific version of CUDA to # work. The cuda12 jaxlib binaries only works with CUDA 12.2, and cuda11 # jaxlib binaries only works with CUDA 11.8. This is why we need to find a # binary that matches the provided cudaVersion. gpuSrcVersionString = "cuda${cudaVersion}-${pythonVersion}"; # Find new releases at https://storage.googleapis.com/jax-releases # When upgrading, you can get these hashes from prefetch.sh. See # https://github.com/google/jax/issues/12879 as to why this specific URL is the correct index. gpuSrcs = { "cuda12.2-3.9" = fetchurl { url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp39-cp39-manylinux2014_x86_64.whl"; hash = "sha256-d8LIl22gIvmWfoyKfXKElZJXicPQIZxdS4HumhwQGCw="; }; "cuda12.2-3.10" = fetchurl { url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp310-cp310-manylinux2014_x86_64.whl"; hash = "sha256-PXtWv+UEcMWF8LhWe6Z1UGkf14PG3dkJ0Iop0LiimnQ="; }; "cuda12.2-3.11" = fetchurl { url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp311-cp311-manylinux2014_x86_64.whl"; hash = "sha256-QO2WSOzmJ48VaCha596mELiOfPsAGLpGctmdzcCHE/o="; }; "cuda12.2-3.12" = fetchurl { url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp312-cp312-manylinux2014_x86_64.whl"; hash = "sha256-ixWMaIChy4Ammsn23/3cCoala0lFibuUxyUr3tjfFKU="; }; }; in buildPythonPackage { pname = "jaxlib"; inherit version; format = "wheel"; disabled = !( pythonVersion == "3.9" || pythonVersion == "3.10" || pythonVersion == "3.11" || pythonVersion == "3.12" ); # See https://discourse.nixos.org/t/ofborg-does-not-respect-meta-platforms/27019/6. src = if !cudaSupport then (cpuSrcs."${pythonVersion}-${stdenv.hostPlatform.system}" or (throw "jaxlib-bin is not supported on ${stdenv.hostPlatform.system}") ) else gpuSrcs."${gpuSrcVersionString}"; # Prebuilt wheels are dynamically linked against things that nix can't find. # Run `autoPatchelfHook` to automagically fix them. nativeBuildInputs = lib.optionals stdenv.isLinux [ autoPatchelfHook ] ++ lib.optionals cudaSupport [ autoAddDriverRunpath ]; # Dynamic link dependencies buildInputs = [ stdenv.cc.cc.lib ]; # jaxlib contains shared libraries that open other shared libraries via dlopen # and these implicit dependencies are not recognized by ldd or # autoPatchelfHook. That means we need to sneak them into rpath. This step # must be done after autoPatchelfHook and the automatic stripping of # artifacts. autoPatchelfHook runs in postFixup and auto-stripping runs in the # patchPhase. preInstallCheck = lib.optional cudaSupport '' shopt -s globstar for file in $out/**/*.so; do patchelf --add-rpath "${cudaLibPath}" "$file" done ''; propagatedBuildInputs = [ absl-py flatbuffers ml-dtypes scipy ]; # jaxlib looks for ptxas at runtime, eg when running `jax.random.PRNGKey(0)`. # Linking into $out is the least bad solution. See # * https://github.com/NixOS/nixpkgs/pull/164176#discussion_r828801621 # * https://github.com/NixOS/nixpkgs/pull/288829#discussion_r1493852211 # for more info. postInstall = lib.optional cudaSupport '' mkdir -p $out/${python.sitePackages}/jaxlib/cuda/bin ln -s ${lib.getExe' cudaPackages.cuda_nvcc "ptxas"} $out/${python.sitePackages}/jaxlib/cuda/bin/ptxas ''; inherit (jaxlib-build) pythonImportsCheck; meta = with lib; { description = "Prebuilt jaxlib backend from PyPi"; homepage = "https://github.com/google/jax"; sourceProvenance = with sourceTypes; [ binaryNativeCode ]; license = licenses.asl20; maintainers = with maintainers; [ samuela ]; platforms = [ "aarch64-darwin" "x86_64-linux" "x86_64-darwin" ]; broken = !(cudaSupport -> lib.versionAtLeast cudaVersion "11.1") || !(cudaSupport -> lib.versionAtLeast cudaPackages.cudnn.version "8.2") || !(cudaSupport -> stdenv.isLinux) || !(cudaSupport -> (gpuSrcs ? "cuda${cudaVersion}-${pythonVersion}")) # Fails at pythonImportsCheckPhase: # ...-python-imports-check-hook.sh/nix-support/setup-hook: line 10: 28017 Illegal instruction: 4 # /nix/store/5qpssbvkzfh73xih07xgmpkj5r565975-python3-3.11.9/bin/python3.11 -c # 'import os; import importlib; list(map(lambda mod: importlib.import_module(mod), os.environ["pythonImportsCheck"].split()))' || (stdenv.isDarwin && stdenv.isx86_64); }; }