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

   1 { lib
   2 , buildPythonPackage
   3 , pythonOlder
   4 , fetchFromGitHub
   5 , pytest-xdist
   6 , pytestCheckHook
   7 , setuptools-scm
   8 , fastprogress
   9 , jax
  10 , jaxlib
  11 , jaxopt
  12 , optax
  13 , typing-extensions
  14 }:
  15
  16 buildPythonPackage rec {
  17   pname = "blackjax";
  18   version = "1.1.1";
  19   pyproject = true;
  20
  21   disabled = pythonOlder "3.9";
  22
  23   src = fetchFromGitHub {
  24     owner = "blackjax-devs";
  25     repo = "blackjax";
  26     rev = "refs/tags/${version}";
  27     hash = "sha256-6+ElY1F8oRCtWT4a/LIG6hYMthlq5mDx2baKAc6zIns=";
  28   };
  29
  30   nativeBuildInputs = [ setuptools-scm ];
  31
  32   propagatedBuildInputs = [
  33     fastprogress
  34     jax
  35     jaxlib
  36     jaxopt
  37     optax
  38     typing-extensions
  39   ];
  40
  41   nativeCheckInputs = [
  42     pytestCheckHook
  43     pytest-xdist
  44   ];
  45   disabledTestPaths = [ "tests/test_benchmarks.py" ];
  46   disabledTests = [
  47     # too slow
  48     "test_adaptive_tempered_smc"
  49   ];
  50
  51   pythonImportsCheck = [
  52     "blackjax"
  53   ];
  54
  55   meta = with lib; {
  56     homepage = "https://blackjax-devs.github.io/blackjax";
  57     description = "Sampling library designed for ease of use, speed and modularity";
  58     changelog = "https://github.com/blackjax-devs/blackjax/releases/tag/${version}";
  59     license = licenses.asl20;
  60     maintainers = with maintainers; [ bcdarwin ];
  61   };
  62 }