From 0df0b7b642ef096a64383bdbc8155e85adf11ba7 Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Wed, 8 Apr 2026 06:17:06 +0000
Subject: [PATCH] python3Packages.orbax-checkpoint: 0.11.33 -> 0.11.39

Diff: https://github.com/google/orbax/compare/v0.11.33...v0.11.39

Changelog:
https://github.com/google/orbax/blob/v0.11.39/checkpoint/CHANGELOG.md
---
 .../orbax-checkpoint/default.nix              | 34 ++++++++++++++-----
 1 file changed, 26 insertions(+), 8 deletions(-)

diff --git a/pkgs/development/python-modules/orbax-checkpoint/default.nix b/pkgs/development/python-modules/orbax-checkpoint/default.nix
index e101000ed58b..cf58f457ff23 100644
--- a/pkgs/development/python-modules/orbax-checkpoint/default.nix
+++ b/pkgs/development/python-modules/orbax-checkpoint/default.nix
@@ -31,7 +31,6 @@
   mock,
   optax,
   portpicker,
-  pytest-xdist,
   pytestCheckHook,
   safetensors,
   torch,
@@ -39,14 +38,15 @@
 
 buildPythonPackage (finalAttrs: {
   pname = "orbax-checkpoint";
-  version = "0.11.33";
+  version = "0.11.39";
   pyproject = true;
+  __structuredAttrs = true;
 
   src = fetchFromGitHub {
     owner = "google";
     repo = "orbax";
     tag = "v${finalAttrs.version}";
-    hash = "sha256-ibHV+MwQvlh2USeDVAUWEJxCS1MGuPZRnZZTBWFO7UQ=";
+    hash = "sha256-KrehggcpKqMd51SdEo3uzYyvH8M15tmECHzvGLBhD/4=";
   };
 
   sourceRoot = "${finalAttrs.src.name}/checkpoint";
@@ -80,15 +80,10 @@ buildPythonPackage (finalAttrs: {
     mock
     optax
     portpicker
-    pytest-xdist
     pytestCheckHook
     safetensors
     torch
   ];
-  pythonImportsCheck = [
-    "orbax"
-    "orbax.checkpoint"
-  ];
 
   disabledTests = [
     # Flaky
@@ -107,6 +102,14 @@ buildPythonPackage (finalAttrs: {
     # AssertionError: False is not true
     "test_register_and_get"
     "test_register_different_modules"
+
+    # IndexError: list index out of range
+    "test_named_sharding"
+
+    # ValueError: cannot reshape array of size 1 into shape (0,2)
+    "test_get_leaf_memory_per_device"
+    "test_number_of_broadcasts"
+    "test_tree_memory_per_device"
   ]
   ++ lib.optionals stdenv.hostPlatform.isDarwin [
     # Probably failing because of a filesystem impurity
@@ -131,8 +134,14 @@ buildPythonPackage (finalAttrs: {
     # Description from first occurrence: Number of processes to use.
     # https://github.com/google/orbax/issues/1580
     "orbax/checkpoint/_src/testing/multiprocess_test.py"
+    "orbax/checkpoint/_src/testing/oss/multiprocess_test.py"
     "orbax/checkpoint/experimental/emergency/"
 
+    # ImportError: cannot import name 'tiering_service_pb2' from
+    # 'orbax.checkpoint.experimental.tiering_service.proto'
+    # (the protobuf module is not generated from the .proto file)
+    "orbax/checkpoint/experimental/tiering_service/server_test.py"
+
     # ValueError: Distributed system is not available; please initialize it via `jax.distributed.initialize()` at the start of your program.
     "orbax/checkpoint/_src/handlers/array_checkpoint_handler_test.py"
 
@@ -148,6 +157,9 @@ buildPythonPackage (finalAttrs: {
     # '/build/absl_testing/DefaultSnapshotTest/runTest/root/path/to/source/data.txt'
     "orbax/checkpoint/_src/path/snapshot/snapshot_test.py"
 
+    # Expects to run on 8 devices
+    "orbax/checkpoint/_src/multihost/multihost_test.py"
+
     # Circular dependency flax
     "orbax/checkpoint/_src/handlers/pytree_checkpoint_handler_test.py"
     "orbax/checkpoint/_src/metadata/empty_values_test.py"
@@ -160,6 +172,12 @@ buildPythonPackage (finalAttrs: {
     "orbax/checkpoint/checkpoint_manager_test.py"
     "orbax/checkpoint/single_host_test.py"
     "orbax/checkpoint/transform_utils_test.py"
+    "orbax/checkpoint/_src/handlers/standard_checkpoint_handler_test.py"
+  ];
+
+  pythonImportsCheck = [
+    "orbax"
+    "orbax.checkpoint"
   ];
 
   meta = {