tjh.dev
nixpkgs mirror (for testing)
github.com/NixOS/nixpkgs
nix
python3Packages.jax: 0.2.21 -> 0.2.24
authored by
Fabian Affolter
and committed by
Jonathan Ringer
ddf89d7d
643912c4
+38
-13
+38
-13
pkgs/development/python-modules/jax/default.nix
+38
-13
pkgs/development/python-modules/jax/default.nix
···
1
-
{ buildPythonPackage, fetchFromGitHub, lib
2
-
# propagatedBuildInputs
3
-
, absl-py, numpy, opt-einsum
4
-
# checkInputs
5
-
, jaxlib, pytestCheckHook
1
+
{ lib
2
+
, absl-py
3
+
, buildPythonPackage
4
+
, fetchFromGitHub
5
+
, jaxlib
6
+
, numpy
7
+
, opt-einsum
8
+
, pytestCheckHook
9
+
, pythonOlder
10
+
, scipy
11
+
, typing-extensions
6
12
}:
7
13
8
14
buildPythonPackage rec {
9
15
pname = "jax";
10
-
version = "0.2.21";
16
+
version = "0.2.24";
17
+
format = "setuptools";
11
18
12
-
# Fetching from pypi doesn't allow us to run the test suite. See https://discourse.nixos.org/t/pythonremovetestsdir-hook-being-run-before-checkphase/14612/3.
19
+
disabled = pythonOlder "3.7";
20
+
13
21
src = fetchFromGitHub {
14
22
owner = "google";
15
23
repo = pname;
16
24
rev = "jax-v${version}";
17
-
sha256 = "05w157h6jv20k8w2gnmlxbycmzf24lr5v392q0c5v0qcql11q7pn";
25
+
sha256 = "1mmn1m4mprpwqlb1smjfdy3f74zm9p3l9dhhn25x6jrcj2cgc5pi";
18
26
};
19
27
20
28
# jaxlib is _not_ included in propagatedBuildInputs because there are
21
29
# different versions of jaxlib depending on the desired target hardware. The
22
30
# JAX project ships separate wheels for CPU, GPU, and TPU. Currently only the
23
31
# CPU wheel is packaged.
24
-
propagatedBuildInputs = [ absl-py numpy opt-einsum ];
32
+
propagatedBuildInputs = [
33
+
absl-py
34
+
numpy
35
+
opt-einsum
36
+
scipy
37
+
typing-extensions
38
+
];
25
39
26
-
checkInputs = [ jaxlib pytestCheckHook ];
40
+
checkInputs = [
41
+
jaxlib
42
+
pytestCheckHook
43
+
];
44
+
27
45
# NOTE: Don't run the tests in the expiremental directory as they require flax
28
46
# which creates a circular dependency. See https://discourse.nixos.org/t/how-to-nix-ify-python-packages-with-circular-dependencies/14648/2.
29
47
# Not a big deal, this is how the JAX docs suggest running the test suite
30
48
# anyhow.
31
-
pytestFlagsArray = [ "-W ignore::DeprecationWarning" "tests/" ];
49
+
pytestFlagsArray = [
50
+
"-W ignore::DeprecationWarning"
51
+
"tests/"
52
+
];
53
+
54
+
pythonImportsCheck = [
55
+
"jax"
56
+
];
32
57
33
58
meta = with lib; {
34
59
description = "Differentiate, compile, and transform Numpy code";
35
-
homepage = "https://github.com/google/jax";
36
-
license = licenses.asl20;
60
+
homepage = "https://github.com/google/jax";
61
+
license = licenses.asl20;
37
62
maintainers = with maintainers; [ samuela ];
38
63
};
39
64
}