2024-06-05 15:53:02 +00:00
|
|
|
{
|
|
|
|
|
lib,
|
|
|
|
|
buildPythonPackage,
|
|
|
|
|
pythonOlder,
|
|
|
|
|
fetchFromGitHub,
|
|
|
|
|
chex,
|
|
|
|
|
jaxlib,
|
|
|
|
|
numpy,
|
|
|
|
|
tensorflow-probability,
|
|
|
|
|
dm-haiku,
|
|
|
|
|
pytest-xdist,
|
|
|
|
|
pytestCheckHook,
|
2023-01-11 07:51:40 +00:00
|
|
|
}:
|
2022-07-18 16:21:45 +00:00
|
|
|
|
|
|
|
|
buildPythonPackage rec {
|
|
|
|
|
pname = "distrax";
|
2024-01-02 11:29:13 +00:00
|
|
|
version = "0.1.5";
|
|
|
|
|
pyproject = true;
|
2022-07-18 16:21:45 +00:00
|
|
|
|
2024-01-02 11:29:13 +00:00
|
|
|
disabled = pythonOlder "3.9";
|
|
|
|
|
|
|
|
|
|
src = fetchFromGitHub {
|
|
|
|
|
owner = "google-deepmind";
|
|
|
|
|
repo = "distrax";
|
|
|
|
|
rev = "refs/tags/v${version}";
|
|
|
|
|
hash = "sha256-A1aCL/I89Blg9sNmIWQru4QJteUTN6+bhgrEJPmCrM0=";
|
2022-07-18 16:21:45 +00:00
|
|
|
};
|
|
|
|
|
|
|
|
|
|
buildInputs = [
|
|
|
|
|
chex
|
|
|
|
|
jaxlib
|
|
|
|
|
numpy
|
|
|
|
|
tensorflow-probability
|
|
|
|
|
];
|
|
|
|
|
|
2023-02-02 18:25:31 +00:00
|
|
|
nativeCheckInputs = [
|
2022-07-18 16:21:45 +00:00
|
|
|
dm-haiku
|
2024-01-13 08:15:51 +00:00
|
|
|
pytest-xdist
|
2022-07-18 16:21:45 +00:00
|
|
|
pytestCheckHook
|
|
|
|
|
];
|
|
|
|
|
|
2024-06-05 15:53:02 +00:00
|
|
|
pythonImportsCheck = [ "distrax" ];
|
2022-07-18 16:21:45 +00:00
|
|
|
|
2024-01-02 11:29:13 +00:00
|
|
|
disabledTests = [
|
|
|
|
|
# AssertionError on numerical values
|
|
|
|
|
# Reported upstream in https://github.com/google-deepmind/distrax/issues/267
|
|
|
|
|
"test_method_with_input_unnormalized_probs__with_device"
|
|
|
|
|
"test_method_with_input_unnormalized_probs__with_jit"
|
|
|
|
|
"test_method_with_input_unnormalized_probs__without_device"
|
|
|
|
|
"test_method_with_input_unnormalized_probs__without_jit"
|
|
|
|
|
"test_method_with_value_1d"
|
|
|
|
|
"test_nested_distributions__with_device"
|
|
|
|
|
"test_nested_distributions__without_device"
|
|
|
|
|
"test_nested_distributions__with_jit"
|
|
|
|
|
"test_nested_distributions__without_jit"
|
|
|
|
|
"test_stability__with_device"
|
|
|
|
|
"test_stability__with_jit"
|
|
|
|
|
"test_stability__without_device"
|
|
|
|
|
"test_stability__without_jit"
|
|
|
|
|
"test_von_mises_sample_gradient"
|
|
|
|
|
"test_von_mises_sample_moments"
|
|
|
|
|
];
|
|
|
|
|
|
2022-07-18 16:21:45 +00:00
|
|
|
disabledTestPaths = [
|
|
|
|
|
# TypeErrors
|
|
|
|
|
"distrax/_src/bijectors/tfp_compatible_bijector_test.py"
|
|
|
|
|
"distrax/_src/distributions/distribution_from_tfp_test.py"
|
|
|
|
|
"distrax/_src/distributions/laplace_test.py"
|
|
|
|
|
"distrax/_src/distributions/multinomial_test.py"
|
|
|
|
|
"distrax/_src/distributions/mvn_diag_plus_low_rank_test.py"
|
|
|
|
|
"distrax/_src/distributions/mvn_kl_test.py"
|
|
|
|
|
"distrax/_src/distributions/straight_through_test.py"
|
|
|
|
|
"distrax/_src/distributions/tfp_compatible_distribution_test.py"
|
|
|
|
|
"distrax/_src/distributions/transformed_test.py"
|
|
|
|
|
"distrax/_src/distributions/uniform_test.py"
|
|
|
|
|
"distrax/_src/utils/transformations_test.py"
|
|
|
|
|
];
|
|
|
|
|
|
|
|
|
|
meta = with lib; {
|
|
|
|
|
description = "Probability distributions in JAX";
|
|
|
|
|
homepage = "https://github.com/deepmind/distrax";
|
|
|
|
|
license = licenses.asl20;
|
|
|
|
|
maintainers = with maintainers; [ onny ];
|
2024-04-21 15:54:59 +00:00
|
|
|
# Several tests fail with:
|
|
|
|
|
# AssertionError: [Chex] Assertion assert_type failed: Error in type compatibility check
|
|
|
|
|
broken = true;
|
2022-07-18 16:21:45 +00:00
|
|
|
};
|
|
|
|
|
}
|
