From 32d91aa5dec2e12ea47d1a34bc00a7ed4be2cad5 Mon Sep 17 00:00:00 2001
From: Sarah Clark <seclark@nextquestion.net>
Date: Thu, 12 Mar 2026 16:42:21 -0700
Subject: [PATCH 1/2] python3Packages.orbax-checkpoint: 0.11.30 -> 0.11.33

---
 .../python-modules/orbax-checkpoint/default.nix    | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

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 = {

From 4c245cdaaffd119122205fc87a04e0467ac006fe Mon Sep 17 00:00:00 2001
From: Sarah Clark <seclark@nextquestion.net>
Date: Thu, 12 Mar 2026 16:48:34 -0700
Subject: [PATCH 2/2] python3Packages.orbax-checkpoint: migrate to finalAttrs

---
 .../python-modules/orbax-checkpoint/default.nix        | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/pkgs/development/python-modules/orbax-checkpoint/default.nix b/pkgs/development/python-modules/orbax-checkpoint/default.nix
index a37db7f1fa5d..f2f43a81c987 100644
--- a/pkgs/development/python-modules/orbax-checkpoint/default.nix
+++ b/pkgs/development/python-modules/orbax-checkpoint/default.nix
@@ -37,7 +37,7 @@
   torch,
 }:
 
-buildPythonPackage rec {
+buildPythonPackage (finalAttrs: {
   pname = "orbax-checkpoint";
   version = "0.11.33";
   pyproject = true;
@@ -45,11 +45,11 @@ buildPythonPackage rec {
   src = fetchFromGitHub {
     owner = "google";
     repo = "orbax";
-    tag = "v${version}";
+    tag = "v${finalAttrs.version}";
     hash = "sha256-ibHV+MwQvlh2USeDVAUWEJxCS1MGuPZRnZZTBWFO7UQ=";
   };
 
-  sourceRoot = "${src.name}/checkpoint";
+  sourceRoot = "${finalAttrs.src.name}/checkpoint";
 
   build-system = [ flit-core ];
 
@@ -165,8 +165,8 @@ buildPythonPackage rec {
   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/v${version}/checkpoint/CHANGELOG.md";
+    changelog = "https://github.com/google/orbax/blob/v${finalAttrs.version}/checkpoint/CHANGELOG.md";
     license = lib.licenses.asl20;
     maintainers = with lib.maintainers; [ fab ];
   };
-}
+})