tjh.dev
1# For the moment we only support the CPU and GPU backends of jaxlib. The TPU
2# backend will require some additional work. Those wheels are located here:
3# https://storage.googleapis.com/jax-releases/libtpu_releases.html.
4
5# For future reference, the easiest way to test the GPU backend is to run
6# NIX_PATH=.. nix-shell -p python3 python3Packages.jax "python3Packages.jaxlib-bin.override { cudaSupport = true; }"
7# export XLA_FLAGS=--xla_gpu_force_compilation_parallelism=1
8# python -c "from jax.lib import xla_bridge; assert xla_bridge.get_backend().platform == 'gpu'"
9# python -c "from jax import random; random.PRNGKey(0)"
10# python -c "from jax import random; x = random.normal(random.PRNGKey(0), (100, 100)); x @ x"
11# There's no convenient way to test the GPU backend in the derivation since the
12# nix build environment blocks access to the GPU. See also:
13# * https://github.com/google/jax/issues/971#issuecomment-508216439
14# * https://github.com/google/jax/issues/5723#issuecomment-913038780
15
16{ absl-py
17, addOpenGLRunpath
18, autoPatchelfHook
19, buildPythonPackage
20, config
21, cudnn ? cudaPackages.cudnn
22, fetchurl
23, flatbuffers
24, isPy39
25, lib
26, python
27, scipy
28, stdenv
29 # Options:
30, cudaSupport ? config.cudaSupport or false
31, cudaPackages ? {}
32}:
33
34let
35 inherit (cudaPackages) cudatoolkit cudnn;
36in
37
38assert cudaSupport -> lib.versionAtLeast cudatoolkit.version "11.1";
39assert cudaSupport -> lib.versionAtLeast cudnn.version "8.2";
40
41let
42 version = "0.4.4";
43
44 pythonVersion = python.pythonVersion;
45
46 # Find new releases at https://storage.googleapis.com/jax-releases/jax_releases.html.
47 # When upgrading, you can get these hashes from prefetch.sh. See
48 # https://github.com/google/jax/issues/12879 as to why this specific URL is
49 # the correct index.
50 cpuSrcs = {
51 "x86_64-linux" = fetchurl {
52 url = "https://storage.googleapis.com/jax-releases/nocuda/jaxlib-${version}-cp310-cp310-manylinux2014_x86_64.whl";
53 hash = "sha256-4VT909AB+ti5HzQvsaZWNY6MS/GItlVEFH9qeZnUuKQ=";
54 };
55 "aarch64-darwin" = fetchurl {
56 url = "https://storage.googleapis.com/jax-releases/mac/jaxlib-${version}-cp310-cp310-macosx_11_0_arm64.whl";
57 hash = "sha256-wuOmoCeTldslSa0MommQeTe+RYKhUMam1ZXrgSov+8U=";
58 };
59 "x86_64-darwin" = fetchurl {
60 url = "https://storage.googleapis.com/jax-releases/mac/jaxlib-${version}-cp310-cp310-macosx_10_14_x86_64.whl";
61 hash = "sha256-arfiTw8yafJwjRwJhKby2O7y3+4ksh3PjaKW9JgJ1ok=";
62 };
63 };
64
65 gpuSrc = fetchurl {
66 url = "https://storage.googleapis.com/jax-releases/cuda11/jaxlib-${version}+cuda11.cudnn82-cp310-cp310-manylinux2014_x86_64.whl";
67 hash = "sha256-bJ62DdzuPSV311ZI2R/LJQ3fOkDibtz2+8wDKw31FLk=";
68 };
69in
70buildPythonPackage rec {
71 pname = "jaxlib";
72 inherit version;
73 format = "wheel";
74
75 # At the time of writing (2022-10-19), there are releases for <=3.10.
76 # Supporting all of them is a pain, so we focus on 3.10, the current nixpkgs
77 # python version.
78 disabled = !(pythonVersion == "3.10");
79
80 # See https://discourse.nixos.org/t/ofborg-does-not-respect-meta-platforms/27019/6.
81 src =
82 if !cudaSupport then
83 (
84 cpuSrcs."${stdenv.hostPlatform.system}"
85 or (throw "jaxlib-bin is not supported on ${stdenv.hostPlatform.system}")
86 ) else gpuSrc;
87
88 # Prebuilt wheels are dynamically linked against things that nix can't find.
89 # Run `autoPatchelfHook` to automagically fix them.
90 nativeBuildInputs = lib.optionals cudaSupport [ autoPatchelfHook addOpenGLRunpath ];
91 # Dynamic link dependencies
92 buildInputs = [ stdenv.cc.cc ];
93
94 # jaxlib contains shared libraries that open other shared libraries via dlopen
95 # and these implicit dependencies are not recognized by ldd or
96 # autoPatchelfHook. That means we need to sneak them into rpath. This step
97 # must be done after autoPatchelfHook and the automatic stripping of
98 # artifacts. autoPatchelfHook runs in postFixup and auto-stripping runs in the
99 # patchPhase. Dependencies:
100 # * libcudart.so.11.0 -> cudatoolkit_11.lib
101 # * libcublas.so.11 -> cudatoolkit_11
102 # * libcuda.so.1 -> opengl driver in /run/opengl-driver/lib
103 preInstallCheck = lib.optional cudaSupport ''
104 shopt -s globstar
105
106 addOpenGLRunpath $out/**/*.so
107
108 for file in $out/**/*.so; do
109 rpath=$(patchelf --print-rpath $file)
110 # For some reason `makeLibraryPath` on `cudatoolkit_11` maps to
111 # <cudatoolkit_11.lib>/lib which is different from <cudatoolkit_11>/lib.
112 patchelf --set-rpath "$rpath:${cudatoolkit}/lib:${lib.makeLibraryPath [ cudatoolkit.lib cudnn ]}" $file
113 done
114 '';
115
116 propagatedBuildInputs = [ absl-py flatbuffers scipy ];
117
118 # Note that cudatoolkit is snecessary since jaxlib looks for "ptxas" in $PATH.
119 # See https://github.com/NixOS/nixpkgs/pull/164176#discussion_r828801621 for
120 # more info.
121 postInstall = lib.optional cudaSupport ''
122 mkdir -p $out/bin
123 ln -s ${cudatoolkit}/bin/ptxas $out/bin/ptxas
124 '';
125
126 pythonImportsCheck = [ "jaxlib" ];
127
128 meta = with lib; {
129 description = "XLA library for JAX";
130 homepage = "https://github.com/google/jax";
131 sourceProvenance = with sourceTypes; [ binaryNativeCode ];
132 license = licenses.asl20;
133 maintainers = with maintainers; [ samuela ];
134 platforms = [ "aarch64-darwin" "x86_64-linux" "x86_64-darwin" ];
135 };
136}