diff options
Diffstat (limited to 'pkgs/development/python-modules/flax/default.nix')
-rw-r--r-- | pkgs/development/python-modules/flax/default.nix | 34 |
1 files changed, 21 insertions, 13 deletions
diff --git a/pkgs/development/python-modules/flax/default.nix b/pkgs/development/python-modules/flax/default.nix index 7c443368bd41f..955177b9bbca3 100644 --- a/pkgs/development/python-modules/flax/default.nix +++ b/pkgs/development/python-modules/flax/default.nix @@ -1,32 +1,40 @@ { lib, buildPythonPackage, - cloudpickle, - einops, + pythonOlder, fetchFromGitHub, - jax, + + # build-system jaxlib, - keras, - matplotlib, + pythonRelaxDepsHook, + setuptools-scm, + + # dependencies + jax, msgpack, numpy, optax, orbax-checkpoint, - pytest-xdist, - pytestCheckHook, - pythonOlder, - pythonRelaxDepsHook, pyyaml, rich, - setuptools-scm, - tensorflow, tensorstore, typing-extensions, + + # checks + cloudpickle, + einops, + keras, + pytest-xdist, + pytestCheckHook, + tensorflow, + + # optional-dependencies + matplotlib, }: buildPythonPackage rec { pname = "flax"; - version = "0.8.4"; + version = "0.8.5"; pyproject = true; disabled = pythonOlder "3.9"; @@ -35,7 +43,7 @@ buildPythonPackage rec { owner = "google"; repo = "flax"; rev = "refs/tags/v${version}"; - hash = "sha256-ZwqKZdJ9LOfWTav5nE9xMsMw/DbryqQUuu5fqeugBzY="; + hash = "sha256-6WOFq0758gtNdrlWqSQBlKmWVIGe5e4PAaGrvHoGjr0="; }; build-system = [ |