From 035ae46b0f47e4e724a94a7efa407558817d9d24 Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Thu, 27 Nov 2025 14:35:17 +0000
Subject: [PATCH] python3Packages.orbax-checkpoint: 0.11.28 -> 0.11.30

Diff: https://github.com/google/orbax/compare/v0.11.28...v0.11.30

Changelog: https://github.com/google/orbax/blob/v0.11.30/checkpoint/CHANGELOG.md
---
 .../orbax-checkpoint/default.nix              | 30 +++++++++++++++----
 1 file changed, 24 insertions(+), 6 deletions(-)

diff --git a/pkgs/development/python-modules/orbax-checkpoint/default.nix b/pkgs/development/python-modules/orbax-checkpoint/default.nix
index 33baac38cd79..e39356a36402 100644
--- a/pkgs/development/python-modules/orbax-checkpoint/default.nix
+++ b/pkgs/development/python-modules/orbax-checkpoint/default.nix
@@ -33,28 +33,25 @@
   pytest-xdist,
   pytestCheckHook,
   safetensors,
+  torch,
 }:
 
 buildPythonPackage rec {
   pname = "orbax-checkpoint";
-  version = "0.11.28";
+  version = "0.11.30";
   pyproject = true;
 
   src = fetchFromGitHub {
     owner = "google";
     repo = "orbax";
     tag = "v${version}";
-    hash = "sha256-a7E60fZRmEXTA220mwr7EDMUc+zYbW7wG40vY7NeAOM=";
+    hash = "sha256-y8l0AVGt2t5zLX+x+yuWHsEDy68agpXIkrew+zfYGXU=";
   };
 
   sourceRoot = "${src.name}/checkpoint";
 
   build-system = [ flit-core ];
 
-  pythonRelaxDeps = [
-    "jax"
-  ];
-
   dependencies = [
     absl-py
     aiofiles
@@ -82,6 +79,7 @@ buildPythonPackage rec {
     pytest-xdist
     pytestCheckHook
     safetensors
+    torch
   ];
 
   pythonImportsCheck = [
@@ -115,6 +113,26 @@ buildPythonPackage rec {
   ];
 
   disabledTestPaths = [
+    # import file mismatch:
+    # imported module 'sharding_test' has this __file__ attribute:
+    #   /build/source/checkpoint/orbax/checkpoint/_src/arrays/sharding_test.py
+    # which is not the same as the test file we want to collect:
+    #   /build/source/checkpoint/orbax/checkpoint/_src/metadata/sharding_test.py
+    "orbax/checkpoint/_src/metadata/sharding_test.py"
+
+    # Circular dependency with clu
+    "orbax/checkpoint/_src/testing/benchmarks/array_handler_benchmark_test.py"
+    "orbax/checkpoint/_src/testing/benchmarks/checkpoint_manager_benchmark_test.py"
+    "orbax/checkpoint/_src/testing/benchmarks/checkpoint_manager_perf_benchmark_test.py"
+    "orbax/checkpoint/_src/testing/benchmarks/checkpoint_policy_benchmark_test.py"
+    "orbax/checkpoint/_src/testing/benchmarks/core/config_parsing_test.py"
+    "orbax/checkpoint/_src/testing/benchmarks/core/core_test.py"
+    "orbax/checkpoint/_src/testing/benchmarks/core/metric_test.py"
+    "orbax/checkpoint/_src/testing/benchmarks/emergency_checkpoint_manager_benchmark_test.py"
+    "orbax/checkpoint/_src/testing/benchmarks/multihost_dispatchers_benchmark_test.py"
+    "orbax/checkpoint/_src/testing/benchmarks/pytree_checkpoint_benchmark_test.py"
+    "orbax/checkpoint/_src/testing/benchmarks/single_replica_benchmark_test.py"
+
     # E   absl.flags._exceptions.DuplicateFlagError: The flag 'num_processes' is defined twice.
     # First from multiprocess_test, Second from orbax.checkpoint._src.testing.multiprocess_test.
     # Description from first occurrence: Number of processes to use.