diff --git a/pkgs/development/python-modules/orbax-checkpoint/default.nix b/pkgs/development/python-modules/orbax-checkpoint/default.nix
index cfe81127811a..e3140c40c4f8 100644
--- a/pkgs/development/python-modules/orbax-checkpoint/default.nix
+++ b/pkgs/development/python-modules/orbax-checkpoint/default.nix
@@ -2,10 +2,13 @@
   lib,
   absl-py,
   buildPythonPackage,
-  cached-property,
+
+  # build-system
+  flit-core,
+
+  # dependencies
   etils,
   fetchPypi,
-  flit-core,
   importlib-resources,
   jax,
   jaxlib,
@@ -13,17 +16,21 @@
   nest-asyncio,
   numpy,
   protobuf,
-  pytest-xdist,
-  pytestCheckHook,
-  pythonOlder,
   pyyaml,
   tensorstore,
   typing-extensions,
+
+  # checks
+  google-cloud-logging,
+  mock,
+  pytest-xdist,
+  pytestCheckHook,
+  pythonOlder,
 }:
 
 buildPythonPackage rec {
   pname = "orbax-checkpoint";
-  version = "0.5.17";
+  version = "0.5.20";
   pyproject = true;
 
   disabled = pythonOlder "3.9";
@@ -31,14 +38,13 @@ buildPythonPackage rec {
   src = fetchPypi {
     pname = "orbax_checkpoint";
     inherit version;
-    hash = "sha256-cFV0oLQdk1sXMS/jaYjnLaHzPVfZdzKTeneoTAJ5P5Q=";
+    hash = "sha256-V91BdeaYqMSVeZGrfmwZ17OoeSrnByuc0rJnzls0iE0=";
   };
 
   build-system = [ flit-core ];
 
   dependencies = [
     absl-py
-    cached-property
     etils
     importlib-resources
     jax
@@ -53,6 +59,8 @@ buildPythonPackage rec {
   ];
 
   nativeCheckInputs = [
+    google-cloud-logging
+    mock
     pytest-xdist
     pytestCheckHook
   ];
@@ -68,11 +76,11 @@ buildPythonPackage rec {
     "orbax/checkpoint/utils_test.py"
   ];
 
-  meta = with lib; {
+  meta = {
     description = "Orbax provides common utility libraries for JAX users";
     homepage = "https://github.com/google/orbax/tree/main/checkpoint";
     changelog = "https://github.com/google/orbax/blob/${version}/CHANGELOG.md";
-    license = licenses.asl20;
-    maintainers = with maintainers; [ fab ];
+    license = lib.licenses.asl20;
+    maintainers = with lib.maintainers; [ fab ];
   };
 }