python312Packages.orbax-checkpoint: 0.5.20 -> 0.6.3, python312Packages.tensorstore: 0.1.53 -> 0.1.65 (#341153)

pkgs/development/python-modules/orbax-checkpoint/default.nix

@@ -1,14 +1,16 @@
 {
   lib,
+  stdenv,
   absl-py,
   buildPythonPackage,
+  fetchFromGitHub,
 
   # build-system
   flit-core,
 
   # dependencies
   etils,
-  fetchPypi,
+  humanize,
   importlib-resources,
   jax,
   jaxlib,
@@ -20,32 +22,34 @@
   tensorstore,
   typing-extensions,
 
-  # checks
+  # tests
+  chex,
   google-cloud-logging,
   mock,
   pytest-xdist,
   pytestCheckHook,
-  pythonOlder,
 }:
 
 buildPythonPackage rec {
   pname = "orbax-checkpoint";
-  version = "0.5.20";
+  version = "0.6.3";
   pyproject = true;
 
-  disabled = pythonOlder "3.9";
-
-  src = fetchPypi {
-    pname = "orbax_checkpoint";
-    inherit version;
-    hash = "sha256-V91BdeaYqMSVeZGrfmwZ17OoeSrnByuc0rJnzls0iE0=";
+  src = fetchFromGitHub {
+    owner = "google";
+    repo = "orbax";
+    rev = "refs/tags/v${version}";
+    hash = "sha256-yqccXQGEvxYnWP8rWmRLjjB0pkSEeXBrRvJIwVowUx0=";
   };
 
+  sourceRoot = "${src.name}/checkpoint";
+
   build-system = [ flit-core ];
 
   dependencies = [
     absl-py
     etils
+    humanize
     importlib-resources
     jax
     jaxlib
@@ -59,6 +63,7 @@ buildPythonPackage rec {
   ];
 
   nativeCheckInputs = [
+    chex
     google-cloud-logging
     mock
     pytest-xdist
@@ -70,6 +75,13 @@ buildPythonPackage rec {
     "orbax.checkpoint"
   ];
 
+  disabledTests = lib.optionals stdenv.isDarwin [
+    # Probably failing because of a filesystem impurity
+    # self.assertFalse(os.path.exists(dst_dir))
+    # AssertionError: True is not false
+    "test_create_snapshot"
+  ];
+
   disabledTestPaths = [
     # Circular dependency flax
     "orbax/checkpoint/transform_utils_test.py"

pkgs/development/python-modules/tensorstore/default.nix

@@ -16,17 +16,17 @@ let
     "aarch64-darwin" = "macosx_11_0_arm64";
   };
   hashes = {
-    "310-x86_64-linux" = "sha256-1b6w9wgT6fffTTpJ3MxdPSrFD7Xaby6prQYFljVn4x4=";
-    "311-x86_64-linux" = "sha256-8+HlzaxH30gB5N+ZKR0Oq+yswhq5gjiSF9jVsg8U22E=";
-    "312-x86_64-linux" = "sha256-e8iEQzB4D3RSXgrcPC4me/vsFKoXf1QFNZfQ7968zQE=";
-    "310-aarch64-darwin" = "sha256-2C60yJk/Pbx2woV7hzEmWGzNKWWnySDfTPm247PWIRA=";
-    "311-aarch64-darwin" = "sha256-rdLB7l/8ZYjV589qKtORiyu1rC7W30wzrsz1uihNRpk=";
-    "312-aarch64-darwin" = "sha256-DpbYMIbqceQeiL7PYwnvn9jLtv8EmfHXmxvPfZCw914=";
+    "310-x86_64-linux" = "sha256-oB68FjYzmRARWpbajQuLpAzWwg9CCji4tLZRFCsztjk=";
+    "311-x86_64-linux" = "sha256-kGEecBu7b3TFGUIRirI9q2W3nipiQwsh/1OB92RqDB4=";
+    "312-x86_64-linux" = "sha256-Vw8sT5kahSN20BQs3MOYesSUZqk4CuvfZR1z5nAO7g8=";
+    "310-aarch64-darwin" = "sha256-2vuVxmJMx/GeaHgzUS6rRdysQFHreVzZ5IT5YSDUJro=";
+    "311-aarch64-darwin" = "sha256-0xRVDSDE9upz2yU7mzpa3Y6l6M5FWOMAPKWBC8eY3Eo=";
+    "312-aarch64-darwin" = "sha256-i2TmLOl2aHD5iyzF6YpjbHKFmBGPx5ixPYyNKKQfRNM=";
   };
 in
 buildPythonPackage rec {
   pname = "tensorstore";
-  version = "0.1.53";
+  version = "0.1.65";
   format = "wheel";
 
   # The source build involves some wonky Bazel stuff.
@@ -44,19 +44,19 @@ buildPythonPackage rec {
 
   nativeBuildInputs = lib.optionals stdenv.isLinux [ autoPatchelfHook ];
 
-  propagatedBuildInputs = [
+  dependencies = [
     ml-dtypes
     numpy
   ];
 
   pythonImportsCheck = [ "tensorstore" ];
 
-  meta = with lib; {
+  meta = {
     description = "Library for reading and writing large multi-dimensional arrays";
     homepage = "https://google.github.io/tensorstore";
     changelog = "https://github.com/google/tensorstore/releases/tag/v${version}";
-    license = licenses.asl20;
-    sourceProvenance = [ sourceTypes.binaryNativeCode ];
-    maintainers = with maintainers; [ samuela ];
+    license = lib.licenses.asl20;
+    sourceProvenance = [ lib.sourceTypes.binaryNativeCode ];
+    maintainers = with lib.maintainers; [ samuela ];
   };
 }