Merge pull request #245792 from samuela/samuela/flax

python3Packages.flax: fix build

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

@@ -10,7 +10,9 @@
 , optax
 , pytest-xdist
 , pytestCheckHook
+, pythonRelaxDepsHook
 , tensorflow
+, tensorstore
 , fetchpatch
 , rich
 }:
@@ -26,7 +28,7 @@ buildPythonPackage rec {
     hash = "sha256-Vv68BK83gTIKj0r9x+twdhqmRYziD0vxQCdHkYSeTak=";
   };
 
-  buildInputs = [ jaxlib ];
+  nativeBuildInputs = [ jaxlib pythonRelaxDepsHook ];
 
   propagatedBuildInputs = [
     jax
@@ -35,8 +37,12 @@ buildPythonPackage rec {
     numpy
     optax
     rich
+    tensorstore
   ];
 
+  # See https://github.com/google/flax/pull/2882.
+  pythonRemoveDeps = [ "orbax" ];
+
   pythonImportsCheck = [
     "flax"
   ];
@@ -64,6 +70,12 @@ buildPythonPackage rec {
     # `tensorflow_datasets`, `vocabulary`) so the benefits of trying to run them
     # would be limited anyway.
     "examples/*"
+
+    # See https://github.com/google/flax/issues/3232.
+    "tests/jax_utils_test.py"
+
+    # Requires orbax which is not packaged as of 2023-07-27.
+    "tests/checkpoints_test.py"
   ];
 
   disabledTests = [
@@ -88,7 +100,5 @@ buildPythonPackage rec {
     changelog = "https://github.com/google/flax/releases/tag/v${version}";
     license = licenses.asl20;
     maintainers = with maintainers; [ ndl ];
-    # Requires orbax which is not available
-    broken = true;
   };
 }

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

@@ -0,0 +1,53 @@
+{ autoPatchelfHook
+, buildPythonPackage
+, fetchPypi
+, lib
+, numpy
+, python
+, stdenv
+}:
+
+let
+  pythonVersionNoDot = builtins.replaceStrings [ "." ] [ "" ] python.pythonVersion;
+  systemToPlatform = {
+    "x86_64-linux" = "manylinux_2_17_x86_64.manylinux2014_x86_64";
+    "aarch64-darwin" = "macosx_11_0_arm64";
+  };
+  hashes = {
+    "310-x86_64-linux" = "sha256-Zuy2zBLV950CMbdtpLNpIWqnXHw2jkjrZG48eGtm42w=";
+    "311-x86_64-linux" = "sha256-Bg5j8QB5z8Ju4bEQsZDojJHTJ4UoQF1pkd4ma83Sc/s=";
+    "310-aarch64-darwin" = "sha256-6Tta4ru1TnobFa4FXWz8fm9rAxF0G09Y2Pj/KaQPVnE=";
+    "311-aarch64-darwin" = "sha256-Sb0tv9ZPQJ4n9b0ybpjJWpreQPZvSC5Sd7CXuUwHCn0=";
+  };
+in
+buildPythonPackage rec {
+  pname = "tensorstore";
+  version = "0.1.40";
+  format = "wheel";
+
+  # The source build involves some wonky Bazel stuff.
+  src = fetchPypi {
+    inherit pname version;
+    format = "wheel";
+    python = "cp${pythonVersionNoDot}";
+    abi = "cp${pythonVersionNoDot}";
+    dist = "cp${pythonVersionNoDot}";
+    platform = systemToPlatform.${stdenv.system} or (throw "unsupported system");
+    hash = hashes."${pythonVersionNoDot}-${stdenv.system}" or (throw "unsupported system/python version combination");
+  };
+
+  nativeBuildInputs = lib.optionals stdenv.isLinux [ autoPatchelfHook ];
+
+  propagatedBuildInputs = [ numpy ];
+
+  pythonImportsCheck = [ "tensorstore" ];
+
+  meta = with lib; {
+    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 ];
+  };
+}

pkgs/top-level/python-packages.nix

@@ -12468,6 +12468,8 @@ self: super: with self; {
 
   tensorrt = callPackage ../development/python-modules/tensorrt { };
 
+  tensorstore = callPackage ../development/python-modules/tensorstore { };
+
   tellduslive = callPackage ../development/python-modules/tellduslive { };
 
   termcolor = callPackage ../development/python-modules/termcolor { };