Unverified Commit 878e3aae authored by Samuel Ainsworth's avatar Samuel Ainsworth Committed by GitHub
Browse files

Merge pull request #245792 from samuela/samuela/flax 

python3Packages.flax: fix build
parents f7505dcd ef573449

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

+13 −3
Original line number Diff line number Diff line
@@ -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 → 100644
+53 −0
Original line number Diff line number Diff line 
{ 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

+2 −0
Original line number Diff line number Diff line
@@ -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 { };