pkgs/development/python-modules/tensorflow-probability/default.nix

   1 { lib
   2 , stdenv
   3 , fetchFromGitHub
   4 , bazel_6
   5 , buildBazelPackage
   6 , buildPythonPackage
   7 , cctools
   8 , python
   9 , setuptools
  10 , wheel
  11 , absl-py
  12 , tensorflow
  13 , six
  14 , numpy
  15 , dm-tree
  16 , keras
  17 , decorator
  18 , cloudpickle
  19 , gast
  20 , hypothesis
  21 , scipy
  22 , pandas
  23 , mpmath
  24 , matplotlib
  25 , mock
  26 , pytest
  27 }:
  28
  29 let
  30   version = "0.21.0";
  31   pname = "tensorflow-probability";
  32
  33   # first build all binaries and generate setup.py using bazel
  34   bazel-wheel = buildBazelPackage {
  35     name = "tensorflow_probability-${version}-py2.py3-none-any.whl";
  36     src = fetchFromGitHub {
  37       owner = "tensorflow";
  38       repo = "probability";
  39       rev = "refs/tags/v${version}";
  40       hash = "sha256-DsJd1E5n86xNS7Ci0DXxoUxQ9jH8OwTZq2UuLlQtMUU=";
  41     };
  42     nativeBuildInputs = [
  43       # needed to create the output wheel in installPhase
  44       python
  45       setuptools
  46       wheel
  47       absl-py
  48       tensorflow
  49     ];
  50
  51     bazel = bazel_6;
  52
  53     bazelTargets = [ ":pip_pkg" ];
  54     LIBTOOL = lib.optionalString stdenv.isDarwin "${cctools}/bin/libtool";
  55
  56     fetchAttrs = {
  57       sha256 = "sha256-1iO/eXz1wvSIRTmGuGZDF9VeDVTiWYnjw0Cby4n/6HM=";
  58     };
  59
  60     buildAttrs = {
  61       preBuild = ''
  62         patchShebangs .
  63       '';
  64
  65       installPhase = ''
  66         # work around timestamp issues
  67         # https://github.com/NixOS/nixpkgs/issues/270#issuecomment-467583872
  68         export SOURCE_DATE_EPOCH=315532800
  69
  70         # First build, then move. Otherwise pip_pkg would create the dir $out
  71         # and then put the wheel in that directory. However we want $out to
  72         # point directly to the wheel file.
  73         ./bazel-bin/pip_pkg . --release
  74         mv *.whl "$out"
  75       '';
  76     };
  77   };
  78 in buildPythonPackage {
  79   inherit version pname;
  80   format = "wheel";
  81
  82   src = bazel-wheel;
  83
  84   propagatedBuildInputs = [
  85     tensorflow
  86     six
  87     numpy
  88     decorator
  89     cloudpickle
  90     gast
  91     dm-tree
  92     keras
  93   ];
  94
  95   # Listed here:
  96   # https://github.com/tensorflow/probability/blob/f3777158691787d3658b5e80883fe1a933d48989/testing/dependency_install_lib.sh#L83
  97   nativeCheckInputs = [
  98     hypothesis
  99     pytest
 100     scipy
 101     pandas
 102     mpmath
 103     matplotlib
 104     mock
 105   ];
 106
 107   # Ideally, we run unit tests with pytest, but in checkPhase, only the Bazel-build wheel is available.
 108   # But it seems not guaranteed that running the tests with pytest will even work, see
 109   # https://github.com/tensorflow/probability/blob/c2a10877feb2c4c06a4dc58281e69c37a11315b9/CONTRIBUTING.md?plain=1#L69
 110   # Ideally, tests would be run using Bazel. For now, lets's do a...
 111
 112   # sanity check
 113   pythonImportsCheck = [ "tensorflow_probability" ];
 114
 115   meta = with lib; {
 116     description = "Library for probabilistic reasoning and statistical analysis";
 117     homepage = "https://www.tensorflow.org/probability/";
 118     license = licenses.asl20;
 119     maintainers = with maintainers; [ GaetanLepage ];
 120   };
 121 }