diff --git a/pkgs/development/python-modules/orbax-checkpoint/default.nix b/pkgs/development/python-modules/orbax-checkpoint/default.nix index e29883c9d3aa..a37db7f1fa5d 100644 --- a/pkgs/development/python-modules/orbax-checkpoint/default.nix +++ b/pkgs/development/python-modules/orbax-checkpoint/default.nix @@ -14,7 +14,6 @@ humanize, jax, msgpack, - nest-asyncio, numpy, protobuf, psutil, @@ -22,10 +21,13 @@ simplejson, tensorstore, typing-extensions, + uvloop, # tests chex, + fastapi, google-cloud-logging, + httpx, mock, optax, portpicker, @@ -37,14 +39,14 @@ buildPythonPackage rec { pname = "orbax-checkpoint"; - version = "0.11.30"; + version = "0.11.33"; pyproject = true; src = fetchFromGitHub { owner = "google"; repo = "orbax"; tag = "v${version}"; - hash = "sha256-y8l0AVGt2t5zLX+x+yuWHsEDy68agpXIkrew+zfYGXU="; + hash = "sha256-ibHV+MwQvlh2USeDVAUWEJxCS1MGuPZRnZZTBWFO7UQ="; }; sourceRoot = "${src.name}/checkpoint"; @@ -58,7 +60,6 @@ buildPythonPackage rec { humanize jax msgpack - nest-asyncio numpy protobuf psutil @@ -66,13 +67,16 @@ buildPythonPackage rec { simplejson tensorstore typing-extensions + uvloop ] ++ etils.optional-dependencies.epath ++ etils.optional-dependencies.epy; nativeCheckInputs = [ chex + fastapi google-cloud-logging + httpx mock optax portpicker @@ -81,7 +85,6 @@ buildPythonPackage rec { safetensors torch ]; - pythonImportsCheck = [ "orbax" "orbax.checkpoint" @@ -156,6 +159,7 @@ buildPythonPackage rec { "orbax/checkpoint/checkpoint_manager_test.py" "orbax/checkpoint/single_host_test.py" "orbax/checkpoint/transform_utils_test.py" + "orbax/checkpoint/_src/handlers/pytree_checkpoint_handler_test.py" ]; meta = {