pkgs/development/python-modules/dm-haiku/default.nix

   1 { buildPythonPackage
   2 , fetchFromGitHub
   3 , fetchpatch
   4 , callPackage
   5 , lib
   6 , jmp
   7 , tabulate
   8 , jaxlib
   9 }:
  10
  11 buildPythonPackage rec {
  12   pname = "dm-haiku";
  13   version = "0.0.10";
  14
  15   src = fetchFromGitHub {
  16     owner = "deepmind";
  17     repo = pname;
  18     rev = "refs/tags/v${version}";
  19     hash = "sha256-EZx3o6PgTeFjTwI9Ko9H39EqPSE0yLWWpsdqX6ALlo4=";
  20   };
  21
  22   patches = [
  23     # https://github.com/deepmind/dm-haiku/issues/717
  24     (fetchpatch {
  25       name = "remove-typing-extensions.patch";
  26       url = "https://github.com/deepmind/dm-haiku/commit/c22867db1a3314a382bd2ce36511e2b756dc32a8.patch";
  27       hash = "sha256-SxJc8FrImwMqTJ5OuJ1f4T+HfHgW/sGqXeIqlxEatlE=";
  28     })
  29     # https://github.com/deepmind/dm-haiku/pull/672
  30     (fetchpatch {
  31       name = "fix-find-namespace-packages.patch";
  32       url = "https://github.com/deepmind/dm-haiku/commit/728031721f77d9aaa260bba0eddd9200d107ba5d.patch";
  33       hash = "sha256-qV94TdJnphlnpbq+B0G3KTx5CFGPno+8FvHyu/aZeQE=";
  34     })
  35   ];
  36
  37   outputs = [
  38     "out"
  39     "testsout"
  40   ];
  41
  42   propagatedBuildInputs = [
  43     jaxlib
  44     jmp
  45     tabulate
  46   ];
  47
  48   pythonImportsCheck = [
  49     "haiku"
  50   ];
  51
  52   postInstall = ''
  53     mkdir $testsout
  54     cp -R examples $testsout/examples
  55   '';
  56
  57   # check in passthru.tests.pytest to escape infinite recursion with bsuite
  58   doCheck = false;
  59
  60   passthru.tests = {
  61     pytest = callPackage ./tests.nix { };
  62   };
  63
  64   meta = with lib; {
  65     description = "Haiku is a simple neural network library for JAX developed by some of the authors of Sonnet.";
  66     homepage = "https://github.com/deepmind/dm-haiku";
  67     license = licenses.asl20;
  68     maintainers = with maintainers; [ ndl ];
  69   };
  70 }