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

   1 { lib
   2 , buildPythonPackage
   3 , pythonOlder
   4 , fetchPypi
   5 , jax
   6 , jaxlib
   7 , multipledispatch
   8 , numpy
   9 , tqdm
  10 , funsor
  11 , pytestCheckHook
  12 # TODO: uncomment when tensorflow-probability gets fixed.
  13 # , tensorflow-probability
  14 }:
  15
  16 buildPythonPackage rec {
  17   pname = "numpyro";
  18   version = "0.13.2";
  19   format = "setuptools";
  20
  21   disabled = pythonOlder "3.9";
  22
  23   src = fetchPypi {
  24     inherit version pname;
  25     hash = "sha256-Um8LFVGAlMeOaN9uMwycHJzqEnTaxp8FYXIk+m2VTug=";
  26   };
  27
  28   propagatedBuildInputs = [
  29     jax
  30     jaxlib
  31     multipledispatch
  32     numpy
  33     tqdm
  34   ];
  35
  36   nativeCheckInputs = [
  37     funsor
  38     pytestCheckHook
  39     # TODO: uncomment when tensorflow-probability gets fixed.
  40     # tensorflow-probability
  41   ];
  42
  43   pythonImportsCheck = [
  44     "numpyro"
  45   ];
  46
  47   disabledTests = [
  48     # AssertionError due to tolerance issues
  49     "test_beta_binomial_log_prob"
  50     "test_collapse_beta"
  51     "test_cpu"
  52     "test_gamma_poisson"
  53     "test_gof"
  54     "test_hpdi"
  55     "test_kl_dirichlet_dirichlet"
  56     "test_kl_univariate"
  57     "test_mean_var"
  58     # Tests want to download data
  59     "data_load"
  60     "test_jsb_chorales"
  61     # RuntimeWarning: overflow encountered in cast
  62     "test_zero_inflated_logits_probs_agree"
  63     # NameError: unbound axis name: _provenance
  64     "test_model_transformation"
  65   ];
  66
  67   # TODO: remove when tensorflow-probability gets fixed.
  68   disabledTestPaths = [
  69     "test/test_distributions.py"
  70   ];
  71
  72   meta = with lib; {
  73     description = "Library for probabilistic programming with NumPy";
  74     homepage = "https://num.pyro.ai/";
  75     changelog = "https://github.com/pyro-ppl/numpyro/releases/tag/${version}";
  76     license = licenses.asl20;
  77     maintainers = with maintainers; [ fab ];
  78   };
  79 }