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

   1 { lib
   2 , buildPythonPackage
   3 , pythonOlder
   4 , fetchFromGitHub
   5 , chex
   6 , jaxlib
   7 , numpy
   8 , tensorflow-probability
   9 , dm-haiku
  10 , pytest-xdist
  11 , pytestCheckHook
  12 }:
  13
  14 buildPythonPackage rec {
  15   pname = "distrax";
  16   version = "0.1.5";
  17   pyproject = true;
  18
  19   disabled = pythonOlder "3.9";
  20
  21   src = fetchFromGitHub {
  22     owner = "google-deepmind";
  23     repo = "distrax";
  24     rev = "refs/tags/v${version}";
  25     hash = "sha256-A1aCL/I89Blg9sNmIWQru4QJteUTN6+bhgrEJPmCrM0=";
  26   };
  27
  28   buildInputs = [
  29     chex
  30     jaxlib
  31     numpy
  32     tensorflow-probability
  33   ];
  34
  35   nativeCheckInputs = [
  36     dm-haiku
  37     pytest-xdist
  38     pytestCheckHook
  39   ];
  40
  41   pythonImportsCheck = [
  42     "distrax"
  43   ];
  44
  45   disabledTests = [
  46     # AssertionError on numerical values
  47     # Reported upstream in https://github.com/google-deepmind/distrax/issues/267
  48     "test_method_with_input_unnormalized_probs__with_device"
  49     "test_method_with_input_unnormalized_probs__with_jit"
  50     "test_method_with_input_unnormalized_probs__without_device"
  51     "test_method_with_input_unnormalized_probs__without_jit"
  52     "test_method_with_value_1d"
  53     "test_nested_distributions__with_device"
  54     "test_nested_distributions__without_device"
  55     "test_nested_distributions__with_jit"
  56     "test_nested_distributions__without_jit"
  57     "test_stability__with_device"
  58     "test_stability__with_jit"
  59     "test_stability__without_device"
  60     "test_stability__without_jit"
  61     "test_von_mises_sample_gradient"
  62     "test_von_mises_sample_moments"
  63   ];
  64
  65   disabledTestPaths = [
  66     # TypeErrors
  67     "distrax/_src/bijectors/tfp_compatible_bijector_test.py"
  68     "distrax/_src/distributions/distribution_from_tfp_test.py"
  69     "distrax/_src/distributions/laplace_test.py"
  70     "distrax/_src/distributions/multinomial_test.py"
  71     "distrax/_src/distributions/mvn_diag_plus_low_rank_test.py"
  72     "distrax/_src/distributions/mvn_kl_test.py"
  73     "distrax/_src/distributions/straight_through_test.py"
  74     "distrax/_src/distributions/tfp_compatible_distribution_test.py"
  75     "distrax/_src/distributions/transformed_test.py"
  76     "distrax/_src/distributions/uniform_test.py"
  77     "distrax/_src/utils/transformations_test.py"
  78   ];
  79
  80   meta = with lib; {
  81     description = "Probability distributions in JAX";
  82     homepage = "https://github.com/deepmind/distrax";
  83     license = licenses.asl20;
  84     maintainers = with maintainers; [ onny ];
  85     # Several tests fail with:
  86     # AssertionError: [Chex] Assertion assert_type failed: Error in type compatibility check
  87     broken = true;
  88   };
  89 }