Something went wrong. Try again.
What happens if I push all of nixpkgs
Something went wrong. Try again.
8.3 kB · 217 lines
Nix
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218# 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.
# For future reference, the easiest way to test the GPU backend is to run# NIX_PATH=.. nix-shell -p python3 python3Packages.jax "python3Packages.jaxlib-bin.override { cudaSupport = true; }"# export XLA_FLAGS=--xla_gpu_force_compilation_parallelism=1# python -c "from jax.lib import xla_bridge; assert xla_bridge.get_backend().platform == 'gpu'"# python -c "from jax import random; random.PRNGKey(0)"# python -c "from jax import random; x = random.normal(random.PRNGKey(0), (100, 100)); x @ x"# There's no convenient way to test the GPU backend in the derivation since the# nix build environment blocks access to the GPU. See also:# * https://github.com/google/jax/issues/971#issuecomment-508216439# * https://github.com/google/jax/issues/5723#issuecomment-913038780
{ absl-py, addOpenGLRunpath, autoPatchelfHook, buildPythonPackage, config, fetchPypi, fetchurl, flatbuffers, jaxlib-build, lib, ml-dtypes, python, scipy, stdenv # Options:, cudaSupport ? config.cudaSupport, cudaPackagesGoogle}:
let inherit (cudaPackagesGoogle) cudatoolkit cudnn;
version = "0.4.23";
inherit (python) pythonVersion;
# 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-maN9RzK6/hYIuPRd8n8n5qa/HyPgAf6UD+mlqzZ1/Xc="; }; "3.9-aarch64-darwin" = getSrcFromPypi { platform = "macosx_11_0_arm64"; dist = "cp39"; hash = "sha256-gdb07c12HCfK5VXT2C+9lYKSiIpPgD8sNmd4eG2M6M4="; }; "3.9-x86_64-darwin" = getSrcFromPypi { platform = "macosx_10_14_x86_64"; dist = "cp39"; hash = "sha256-TdU4wEoqEhsDq18MuLEpmKqpU51+xUYp/reZqEDJK1U="; };
"3.10-x86_64-linux" = getSrcFromPypi { platform = "manylinux2014_x86_64"; dist = "cp310"; hash = "sha256-cnX75aSJxoPFUCYD1V5QgyPNovS9lSGqg4PGdPsKsvM="; }; "3.10-aarch64-darwin" = getSrcFromPypi { platform = "macosx_11_0_arm64"; dist = "cp310"; hash = "sha256-H9sbeR4+4XytREYLP0LJphqGkQqHcinTC9NlT0Rj1aA="; }; "3.10-x86_64-darwin" = getSrcFromPypi { platform = "macosx_10_14_x86_64"; dist = "cp310"; hash = "sha256-43VuBgGvdjauWPQtJK9w5GBI/++JvV4FwwO4maIXfDY="; };
"3.11-x86_64-linux" = getSrcFromPypi { platform = "manylinux2014_x86_64"; dist = "cp311"; hash = "sha256-mEdm0wmyHKg4RlA7q9/e1OOu+BfGcPKBCSvLwXfFhJI="; }; "3.11-aarch64-darwin" = getSrcFromPypi { platform = "macosx_11_0_arm64"; dist = "cp311"; hash = "sha256-1kb/m8DODrtXO2drIfpttCLC72oNVszADbSDspllQVs="; }; "3.11-x86_64-darwin" = getSrcFromPypi { platform = "macosx_10_14_x86_64"; dist = "cp311"; hash = "sha256-jhLX4ps+EtU1sku722v51mz2SSamo4/dkdRWX3zFcRE="; };
"3.12-x86_64-linux" = getSrcFromPypi { platform = "manylinux2014_x86_64"; dist = "cp312"; hash = "sha256-oimiuQopgN1oKhbDc7SsRJPnA6JiEI9UieikWR2qpVk="; }; "3.12-aarch64-darwin" = getSrcFromPypi { platform = "macosx_11_0_arm64"; dist = "cp312"; hash = "sha256-J4zaKcx0c0Bgk7w/n6klqDlgY9IqTNINfg6g033LUDk="; }; "3.12-x86_64-darwin" = getSrcFromPypi { platform = "macosx_10_14_x86_64"; dist = "cp312"; hash = "sha256-UFEE/mBitEOVUoijhUfphyy24QfWPZ+FQPsQ0cjY79A="; }; };
# Find new releases at https://storage.googleapis.com/jax-releases/jax_releases.html. # 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 = { "3.9" = fetchurl { url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp39-cp39-manylinux2014_x86_64.whl"; hash = "sha256-our2mSwHPdjVoDAZP+9aNUkJ+vxv1Tq7G5UqA9HvhNI="; }; "3.10" = fetchurl { url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp310-cp310-manylinux2014_x86_64.whl"; hash = "sha256-jkIABnJZnn7A6n9VGs/MldzdDiKwWh0fEvl7Vqn85Kg="; }; "3.11" = fetchurl { url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp311-cp311-manylinux2014_x86_64.whl"; hash = "sha256-dMUcRnHjl8NyUeO3P1x7CNgF0iAHFKIzUtHh+/CNkow="; }; "3.12" = fetchurl { url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp312-cp312-manylinux2014_x86_64.whl"; hash = "sha256-kXJ6bUwX+QybqYPV9Kpwv+lhdoGEFRr4+1T0vfXoWRo="; }; };
inbuildPythonPackage { 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."${pythonVersion}";
# 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 [ addOpenGLRunpath ]; # 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. Dependencies: # * libcudart.so.11.0 -> cudatoolkit_11.lib # * libcublas.so.11 -> cudatoolkit_11 # * libcuda.so.1 -> opengl driver in /run/opengl-driver/lib preInstallCheck = lib.optional cudaSupport '' shopt -s globstar
addOpenGLRunpath $out/**/*.so
for file in $out/**/*.so; do rpath=$(patchelf --print-rpath $file) # For some reason `makeLibraryPath` on `cudatoolkit_11` maps to # <cudatoolkit_11.lib>/lib which is different from <cudatoolkit_11>/lib. patchelf --set-rpath "$rpath:${cudatoolkit}/lib:${lib.makeLibraryPath [ cudatoolkit.lib cudnn ]}" $file done '';
propagatedBuildInputs = [ absl-py flatbuffers ml-dtypes scipy ];
# Note that cudatoolkit is snecessary since jaxlib looks for "ptxas" in $PATH. # See https://github.com/NixOS/nixpkgs/pull/164176#discussion_r828801621 for # more info. postInstall = lib.optional cudaSupport '' mkdir -p $out/bin ln -s ${cudatoolkit}/bin/ptxas $out/bin/ptxas '';
inherit (jaxlib-build) pythonImportsCheck;
meta = with lib; { description = "XLA library for JAX"; 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 -> (cudaPackagesGoogle ? cudatoolkit) && lib.versionAtLeast cudatoolkit.version "11.1") || !(cudaSupport -> (cudaPackagesGoogle ? cudnn) && lib.versionAtLeast cudnn.version "8.2") || !(cudaSupport -> stdenv.isLinux); };}