pyrox.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, fetchPypi
22, fetchurl
23, flatbuffers
24, jaxlib-build
25, lib
26, ml-dtypes
27, python
28, scipy
29, stdenv
30 # Options:
31, cudaSupport ? config.cudaSupport
32, cudaPackages ? {}
33}:
34
35let
36 inherit (cudaPackages) cudatoolkit cudnn;
37in
38
39assert cudaSupport -> lib.versionAtLeast cudatoolkit.version "11.1" && lib.versionAtLeast cudnn.version "8.2" && stdenv.isLinux;
40
41let
42 version = "0.4.20";
43
44 inherit (python) pythonVersion;
45
46 # As of 2023-06-06, google/jax upstream is no longer publishing CPU-only wheels to their GCS bucket. Instead the
47 # official instructions recommend installing CPU-only versions via PyPI.
48 cpuSrcs =
49 let
50 getSrcFromPypi = { platform, dist, hash }: fetchPypi {
51 inherit version platform dist hash;
52 pname = "jaxlib";
53 format = "wheel";
54 # See the `disabled` attr comment below.
55 python = dist;
56 abi = dist;
57 };
58 in
59 {
60 "3.9-x86_64-linux" = getSrcFromPypi {
61 platform = "manylinux2014_x86_64";
62 dist = "cp39";
63 hash = "sha256-eIE+rz5x5BEkO85zncIWE8p/wDPxV8bnVJdHiknS998=";
64 };
65 "3.9-aarch64-darwin" = getSrcFromPypi {
66 platform = "macosx_11_0_arm64";
67 dist = "cp39";
68 hash = "sha256-dxInv8/aQiHsN7DpScuZao2ZyHDjF0AaTqUDA0qqg/M=";
69 };
70 "3.9-x86_64-darwin" = getSrcFromPypi {
71 platform = "macosx_10_14_x86_64";
72 dist = "cp39";
73 hash = "sha256-wva6LkSokEHN+WQLCancVC7YBIxfImPsQpB1LzFcyqM=";
74 };
75
76 "3.10-x86_64-linux" = getSrcFromPypi {
77 platform = "manylinux2014_x86_64";
78 dist = "cp310";
79 hash = "sha256-Yo2TYnkIelyy4vb5+nC/yY8SjV34i/jJvCe/VRQppmo=";
80 };
81 "3.10-aarch64-darwin" = getSrcFromPypi {
82 platform = "macosx_11_0_arm64";
83 dist = "cp310";
84 hash = "sha256-ufA/ACE4s4R/Fiq5SN7T44SVEN1Z5OfkJ/98lKxRFmo=";
85 };
86 "3.10-x86_64-darwin" = getSrcFromPypi {
87 platform = "macosx_10_14_x86_64";
88 dist = "cp310";
89 hash = "sha256-hBSrYQyOGMn0BexRWQKYnJdEYYlzHUWuWGHmjVT10TE=";
90 };
91
92 "3.11-x86_64-linux" = getSrcFromPypi {
93 platform = "manylinux2014_x86_64";
94 dist = "cp311";
95 hash = "sha256-5N0nghTBrsa7d8kt8hZC2ghqlxCNC7U8ApD0PG7DHn8=";
96 };
97 "3.11-aarch64-darwin" = getSrcFromPypi {
98 platform = "macosx_11_0_arm64";
99 dist = "cp311";
100 hash = "sha256-j13Br64cKe0hFh/cMBbOMuTXqauAvSKE+KzEmN7U6RA=";
101 };
102 "3.11-x86_64-darwin" = getSrcFromPypi {
103 platform = "macosx_10_14_x86_64";
104 dist = "cp311";
105 hash = "sha256-nTnyawU4Ngq9VTE6oDuEfR6iJPRy+E/VLt98cU6eW4M=";
106 };
107
108 "3.12-x86_64-linux" = getSrcFromPypi {
109 platform = "manylinux2014_x86_64";
110 dist = "cp312";
111 hash = "sha256-qPMoa7cso7DRBWuCJQoiOEzLPL3m76MPZZMYmZUj400=";
112 };
113 "3.12-aarch64-darwin" = getSrcFromPypi {
114 platform = "macosx_11_0_arm64";
115 dist = "cp312";
116 hash = "sha256-VqTC5egDHaDIvwVa3sAc9Sdtd0CwEFcXjDU/i54h844=";
117 };
118 "3.12-x86_64-darwin" = getSrcFromPypi {
119 platform = "macosx_10_14_x86_64";
120 dist = "cp312";
121 hash = "sha256-1F98Je2rMJJKrksI/EVAsX9n+dOpmDehUeAaMq/BY7o=";
122 };
123 };
124
125 # Find new releases at https://storage.googleapis.com/jax-releases/jax_releases.html.
126 # When upgrading, you can get these hashes from prefetch.sh. See
127 # https://github.com/google/jax/issues/12879 as to why this specific URL is the correct index.
128 gpuSrcs = {
129 "3.9" = fetchurl {
130 url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp39-cp39-manylinux2014_x86_64.whl";
131 hash = "sha256-VM2HuyMnG+hzrsTQEB5KJpqpBXyyp+eV1LVxmY1ZCGU=";
132 };
133 "3.10" = fetchurl {
134 url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp310-cp310-manylinux2014_x86_64.whl";
135 hash = "sha256-TLq3z3T2fjTcO3ESahboKG33mrOpjtj9C92f4d4nJKo=";
136 };
137 "3.11" = fetchurl {
138 url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp311-cp311-manylinux2014_x86_64.whl";
139 hash = "sha256-CUXwyJq0HOo2j3Sw+NguBCnFkDuJpc3wfZUc90yyhOY=";
140 };
141 "3.12" = fetchurl {
142 url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-${version}+cuda12.cudnn89-cp312-cp312-manylinux2014_x86_64.whl";
143 hash = "sha256-bAR8FLtiqufU+rL2a1q9c61CjH1eXxGTNGnDUkHlDBA=";
144 };
145 };
146
147in
148buildPythonPackage {
149 pname = "jaxlib";
150 inherit version;
151 format = "wheel";
152
153 disabled = !(pythonVersion == "3.9" || pythonVersion == "3.10" || pythonVersion == "3.11" || pythonVersion == "3.12");
154
155 # See https://discourse.nixos.org/t/ofborg-does-not-respect-meta-platforms/27019/6.
156 src =
157 if !cudaSupport then
158 (
159 cpuSrcs."${pythonVersion}-${stdenv.hostPlatform.system}"
160 or (throw "jaxlib-bin is not supported on ${stdenv.hostPlatform.system}")
161 ) else gpuSrcs."${pythonVersion}";
162
163 # Prebuilt wheels are dynamically linked against things that nix can't find.
164 # Run `autoPatchelfHook` to automagically fix them.
165 nativeBuildInputs = lib.optionals stdenv.isLinux [ autoPatchelfHook ]
166 ++ lib.optionals cudaSupport [ addOpenGLRunpath ];
167 # Dynamic link dependencies
168 buildInputs = [ stdenv.cc.cc.lib ];
169
170 # jaxlib contains shared libraries that open other shared libraries via dlopen
171 # and these implicit dependencies are not recognized by ldd or
172 # autoPatchelfHook. That means we need to sneak them into rpath. This step
173 # must be done after autoPatchelfHook and the automatic stripping of
174 # artifacts. autoPatchelfHook runs in postFixup and auto-stripping runs in the
175 # patchPhase. Dependencies:
176 # * libcudart.so.11.0 -> cudatoolkit_11.lib
177 # * libcublas.so.11 -> cudatoolkit_11
178 # * libcuda.so.1 -> opengl driver in /run/opengl-driver/lib
179 preInstallCheck = lib.optional cudaSupport ''
180 shopt -s globstar
181
182 addOpenGLRunpath $out/**/*.so
183
184 for file in $out/**/*.so; do
185 rpath=$(patchelf --print-rpath $file)
186 # For some reason `makeLibraryPath` on `cudatoolkit_11` maps to
187 # <cudatoolkit_11.lib>/lib which is different from <cudatoolkit_11>/lib.
188 patchelf --set-rpath "$rpath:${cudatoolkit}/lib:${lib.makeLibraryPath [ cudatoolkit.lib cudnn ]}" $file
189 done
190 '';
191
192 propagatedBuildInputs = [
193 absl-py
194 flatbuffers
195 ml-dtypes
196 scipy
197 ];
198
199 # Note that cudatoolkit is snecessary since jaxlib looks for "ptxas" in $PATH.
200 # See https://github.com/NixOS/nixpkgs/pull/164176#discussion_r828801621 for
201 # more info.
202 postInstall = lib.optional cudaSupport ''
203 mkdir -p $out/bin
204 ln -s ${cudatoolkit}/bin/ptxas $out/bin/ptxas
205 '';
206
207 inherit (jaxlib-build) pythonImportsCheck;
208
209 meta = with lib; {
210 description = "XLA library for JAX";
211 homepage = "https://github.com/google/jax";
212 sourceProvenance = with sourceTypes; [ binaryNativeCode ];
213 license = licenses.asl20;
214 maintainers = with maintainers; [ samuela ];
215 platforms = [ "aarch64-darwin" "x86_64-linux" "x86_64-darwin" ];
216 };
217}