tjh.dev
1{
2 lib,
3 config,
4 stdenv,
5 blas,
6 lapack,
7 buildPythonPackage,
8 fetchFromGitHub,
9 cudaSupport ? config.cudaSupport,
10
11 # build-system
12 setuptools,
13
14 # dependencies
15 jaxlib,
16 ml-dtypes,
17 numpy,
18 opt-einsum,
19 scipy,
20
21 # optional-dependencies
22 jax-cuda12-plugin,
23
24 # tests
25 cloudpickle,
26 hypothesis,
27 matplotlib,
28 pytestCheckHook,
29 pytest-xdist,
30
31 # passthru
32 callPackage,
33 jax,
34 jaxlib-build,
35 jaxlib-bin,
36}:
37
38let
39 usingMKL = blas.implementation == "mkl" || lapack.implementation == "mkl";
40in
41buildPythonPackage rec {
42 pname = "jax";
43 version = "0.6.2";
44 pyproject = true;
45
46 src = fetchFromGitHub {
47 owner = "google";
48 repo = "jax";
49 # google/jax contains tags for jax and jaxlib. Only use jax tags!
50 tag = "jax-v${version}";
51 hash = "sha256-MTgpwpJWxULCiZhDG+MFpOp8ZHoj1ZDmOD05OaGfXhM=";
52 };
53
54 build-system = [ setuptools ];
55
56 # The version is automatically set to ".dev" if this variable is not set.
57 # https://github.com/google/jax/commit/e01f2617b85c5bdffc5ffb60b3d8d8ca9519a1f3
58 JAX_RELEASE = "1";
59
60 dependencies = [
61 jaxlib
62 ml-dtypes
63 numpy
64 opt-einsum
65 scipy
66 ]
67 ++ lib.optionals cudaSupport optional-dependencies.cuda;
68
69 optional-dependencies = rec {
70 cuda = [ jax-cuda12-plugin ];
71 cuda12 = cuda;
72 cuda12_pip = cuda;
73 cuda12_local = cuda;
74 };
75
76 nativeCheckInputs = [
77 cloudpickle
78 hypothesis
79 matplotlib
80 pytestCheckHook
81 pytest-xdist
82 ];
83
84 # high parallelism will result in the tests getting stuck
85 dontUsePytestXdist = true;
86
87 pytestFlags = [
88 "--numprocesses=4"
89 "-Wignore::DeprecationWarning"
90 ];
91
92 # NOTE: Don't run the tests in the experimental directory as they require flax
93 # which creates a circular dependency. See https://discourse.nixos.org/t/how-to-nix-ify-python-packages-with-circular-dependencies/14648/2.
94 # Not a big deal, this is how the JAX docs suggest running the test suite
95 # anyhow.
96 enabledTestPaths = [
97 "tests/"
98 ];
99
100 disabledTestPaths = lib.optionals stdenv.hostPlatform.isDarwin [
101 # SystemError: nanobind::detail::nb_func_error_except(): exception could not be translated!
102 # reported at: https://github.com/jax-ml/jax/issues/26106
103 "tests/pjit_test.py::PJitErrorTest::testAxisResourcesMismatch"
104 "tests/shape_poly_test.py::ShapePolyTest"
105 "tests/tree_util_test.py::TreeTest"
106 ];
107
108 # Prevents `tests/export_back_compat_test.py::CompatTest::test_*` tests from failing on darwin with
109 # PermissionError: [Errno 13] Permission denied: '/tmp/back_compat_testdata/test_*.py'
110 # See https://github.com/google/jax/blob/jaxlib-v0.4.27/jax/_src/internal_test_util/export_back_compat_test_util.py#L240-L241
111 # NOTE: this doesn't seem to be an issue on linux
112 preCheck = lib.optionalString stdenv.hostPlatform.isDarwin ''
113 export TEST_UNDECLARED_OUTPUTS_DIR=$(mktemp -d)
114 '';
115
116 disabledTests = [
117 # Exceeds tolerance when the machine is busy
118 "test_custom_linear_solve_aux"
119 ]
120 ++ lib.optionals usingMKL [
121 # See
122 # * https://github.com/google/jax/issues/9705
123 # * https://discourse.nixos.org/t/getting-different-results-for-the-same-build-on-two-equally-configured-machines/17921
124 # * https://github.com/NixOS/nixpkgs/issues/161960
125 "test_custom_linear_solve_cholesky"
126 "test_custom_root_with_aux"
127 "testEigvalsGrad_shape"
128 ]
129 ++ lib.optionals stdenv.hostPlatform.isAarch64 [
130 # Fails on some hardware due to some numerical error
131 # See https://github.com/google/jax/issues/18535
132 "testQdwhWithOnRankDeficientInput5"
133 ]
134 ++ lib.optionals stdenv.hostPlatform.isDarwin [
135 # SystemError: nanobind::detail::nb_func_error_except(): exception could not be translated!
136 # reported at: https://github.com/jax-ml/jax/issues/26106
137 "testInAxesPyTreePrefixMismatchError"
138 "testInAxesPyTreePrefixMismatchErrorKwargs"
139 "testOutAxesPyTreePrefixMismatchError"
140 "test_tree_map"
141 "test_tree_prefix_error"
142 "test_vjp_rule_inconsistent_pytree_structures_error"
143 "test_vmap_in_axes_tree_prefix_error"
144 "test_vmap_mismatched_axis_sizes_error_message_issue_705"
145 ];
146
147 pythonImportsCheck = [ "jax" ];
148
149 # Test CUDA-enabled jax and jaxlib. Running CUDA-enabled tests is not
150 # currently feasible within the nix build environment so we have to maintain
151 # this script separately. See https://github.com/NixOS/nixpkgs/pull/256230
152 # for a possible remedy to this situation.
153 #
154 # Run these tests with eg
155 #
156 # NIXPKGS_ALLOW_UNFREE=1 nixglhost -- nix run --impure .#python3Packages.jax.passthru.tests.test_cuda_jaxlibBin
157 passthru.tests = {
158 # jaxlib-build is broken as of 2024-12-20
159 # test_cuda_jaxlibSource = callPackage ./test-cuda.nix {
160 # jax = jax.override { jaxlib = jaxlib-build; };
161 # };
162 test_cuda_jaxlibBin = callPackage ./test-cuda.nix {
163 jax = jax.override { jaxlib = jaxlib-bin; };
164 };
165 };
166
167 # updater fails to pick the correct branch
168 passthru.skipBulkUpdate = true;
169
170 meta = {
171 description = "Source-built JAX frontend: differentiate, compile, and transform Numpy code";
172 homepage = "https://github.com/google/jax";
173 license = lib.licenses.asl20;
174 maintainers = with lib.maintainers; [ samuela ];
175 };
176}