Something went wrong. Try again.
What happens if I push all of nixpkgs
Something went wrong. Try again.
4.2 kB · 127 lines
Nix
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128{ lib, buildPythonPackage, fetchFromGitHub, fetchpatch, chex, jaxlib, numpy, tensorflow-probability, dm-haiku, pytest-xdist, pytestCheckHook,}:
buildPythonPackage rec { pname = "distrax"; version = "0.1.5"; pyproject = true;
src = fetchFromGitHub { owner = "google-deepmind"; repo = "distrax"; tag = "v${version}"; hash = "sha256-A1aCL/I89Blg9sNmIWQru4QJteUTN6+bhgrEJPmCrM0="; };
patches = [ # TODO: remove at the next release (already on master) (fetchpatch { name = "fix-jax-0.6.0-compat"; url = "https://github.com/google-deepmind/distrax/commit/c02708ac46518fac00ab2945311e0f2ee32c672c.patch"; hash = "sha256-hFNXKoA1b5I6dzhwTRXp/SnkHv89GI6tYwlnBBHwG78="; }) ];
dependencies = [ chex jaxlib numpy tensorflow-probability ];
nativeCheckInputs = [ dm-haiku pytest-xdist pytestCheckHook ];
pythonImportsCheck = [ "distrax" ];
disabledTests = [ # Flaky: AssertionError: 1 not less than 0.7000000000000001 "test_von_mises_sample_uniform_ks_test"
# Flaky: AssertionError: Not equal to tolerance "test_composite_methods_are_consistent__with_jit"
# NotImplementedError: Primitive 'square' does not have a registered inverse. "test_against_tfp_bijectors_square" "test_log_dets_square__with_device" "test_log_dets_square__without_device" "test_log_dets_square__without_jit"
# 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" ];
disabledTestPaths = [ # Since jax 0.6.0: # TypeError: <lambda>() got an unexpected keyword argument 'accuracy' "distrax/_src/bijectors/lambda_bijector_test.py"
# 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" # https://github.com/google-deepmind/distrax/pull/270 "distrax/_src/distributions/deterministic_test.py" "distrax/_src/distributions/epsilon_greedy_test.py" "distrax/_src/distributions/gamma_test.py" "distrax/_src/distributions/greedy_test.py" "distrax/_src/distributions/gumbel_test.py" "distrax/_src/distributions/logistic_test.py" "distrax/_src/distributions/log_stddev_normal_test.py" "distrax/_src/distributions/mvn_diag_test.py" "distrax/_src/distributions/mvn_full_covariance_test.py" "distrax/_src/distributions/mvn_tri_test.py" "distrax/_src/distributions/one_hot_categorical_test.py" "distrax/_src/distributions/softmax_test.py" "distrax/_src/utils/hmm_test.py" ];
meta = { description = "Probability distributions in JAX"; homepage = "https://github.com/deepmind/distrax"; changelog = "https://github.com/google-deepmind/distrax/releases/tag/v${version}"; license = lib.licenses.asl20; maintainers = with lib.maintainers; [ onny ]; badPlatforms = [ # SystemError: nanobind::detail::nb_func_error_except(): exception could not be translated! lib.systems.inspect.patterns.isDarwin ]; };}