From 46388074bef0736c8b5598d5efc8f50134f6653d Mon Sep 17 00:00:00 2001 From: Samuel Ainsworth Date: Thu, 27 Jul 2023 17:17:47 -0400 Subject: [PATCH 1/3] python3Packages.tensorstore: init at 0.1.40 --- .../python-modules/tensorstore/default.nix | 52 +++++++++++++++++++ pkgs/top-level/python-packages.nix | 2 + 2 files changed, 54 insertions(+) create mode 100644 pkgs/development/python-modules/tensorstore/default.nix diff --git a/pkgs/development/python-modules/tensorstore/default.nix b/pkgs/development/python-modules/tensorstore/default.nix new file mode 100644 index 000000000000..fdb5725372e6 --- /dev/null +++ b/pkgs/development/python-modules/tensorstore/default.nix @@ -0,0 +1,52 @@ +{ 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; + maintainers = with maintainers; [ samuela ]; + }; +} diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index 088b79d86c37..6e3c2f015d51 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -12424,6 +12424,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 { }; From 33e976114d3cfcb2f5b7d442eac3e62228c9f048 Mon Sep 17 00:00:00 2001 From: Samuel Ainsworth Date: Thu, 27 Jul 2023 17:18:07 -0400 Subject: [PATCH 2/3] python3Packages.flax: fix build --- pkgs/development/python-modules/flax/default.nix | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/pkgs/development/python-modules/flax/default.nix b/pkgs/development/python-modules/flax/default.nix index dda4a7ffc053..bea9667a1a18 100644 --- a/pkgs/development/python-modules/flax/default.nix +++ b/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; }; } From ef573449913f7c3c883945aac284b2d036aad2a4 Mon Sep 17 00:00:00 2001 From: Samuel Ainsworth Date: Sun, 30 Jul 2023 21:35:54 -0700 Subject: [PATCH 3/3] python3Packages.tensorstore: specify meta.sourceProvenance Co-authored-by: Robert Scott --- pkgs/development/python-modules/tensorstore/default.nix | 1 + 1 file changed, 1 insertion(+) diff --git a/pkgs/development/python-modules/tensorstore/default.nix b/pkgs/development/python-modules/tensorstore/default.nix index fdb5725372e6..5e8631062394 100644 --- a/pkgs/development/python-modules/tensorstore/default.nix +++ b/pkgs/development/python-modules/tensorstore/default.nix @@ -47,6 +47,7 @@ buildPythonPackage rec { 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 ]; }; }