pyrox.dev
python3Packages.optax: init at unstable-2022-01-05 (#153655)
authored by Alexander Tsvyashchenko and committed by GitHub 12e13654 0c8b2ce6
+55
pkgs/development/python-modules/optax/default.nix
+55
pkgs/development/python-modules/optax/default.nix
···
···
1
+
{ absl-py
2
+
, buildPythonPackage
3
+
, chex
4
+
, dm-haiku
5
+
, fetchFromGitHub
6
+
, jaxlib
7
+
, lib
8
+
, numpy
9
+
, pytestCheckHook
10
+
}:
11
+
12
+
buildPythonPackage rec {
13
+
pname = "optax";
14
+
# As of 2022-01-06, the latest stable version (0.1.0) has broken tests that are fixed
15
+
# in https://github.com/deepmind/optax/commit/d6633365d84eb6f2c0df0c52b630481a349ce562
16
+
version = "unstable-2022-01-05";
17
+
18
+
src = fetchFromGitHub {
19
+
owner = "deepmind";
20
+
repo = pname;
21
+
rev = "5ec5541b3486224b22e950480ff639ceaf5098f7";
22
+
sha256 = "1q8cxc42a5xais2ll1l238cnn3l7w28savhgiz0lg01ilz2ysbli";
23
+
};
24
+
25
+
propagatedBuildInputs = [
26
+
absl-py
27
+
chex
28
+
jaxlib
29
+
numpy
30
+
];
31
+
32
+
checkInputs = [
33
+
dm-haiku
34
+
pytestCheckHook
35
+
];
36
+
37
+
pythonImportsCheck = [
38
+
"optax"
39
+
];
40
+
41
+
disabledTestPaths = [
42
+
# Requires `flax` which depends on `optax` creating circular dependency.
43
+
"optax/_src/equivalence_test.py"
44
+
# Require `tensorflow_datasets` which isn't packaged in `nixpkgs`.
45
+
"examples/datasets_test.py"
46
+
"examples/lookahead_mnist_test.py"
47
+
];
48
+
49
+
meta = with lib; {
50
+
description = "Optax is a gradient processing and optimization library for JAX.";
51
+
homepage = "https://github.com/deepmind/optax";
52
+
license = licenses.asl20;
53
+
maintainers = with maintainers; [ ndl ];
54
+
};
55
+
}
+2
pkgs/top-level/python-packages.nix
+2
pkgs/top-level/python-packages.nix
···
5600
5601
opt-einsum = callPackage ../development/python-modules/opt-einsum { };
5602
5603
+
optax = callPackage ../development/python-modules/optax { };
5604
+
5605
optuna = callPackage ../development/python-modules/optuna { };
5606
5607
opuslib = callPackage ../development/python-modules/opuslib { };