Something went wrong. Try again.
Nix Packages collection & NixOS github.com/NixOS/nixpkgs
Something went wrong. Try again.
12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697{ lib, buildPythonPackage, fetchFromGitHub,
# build-system flit-core,
# dependencies absl-py, chex, jax, jaxlib, numpy, tensorflow-probability,
# tests dm-haiku, mock, pytest-xdist, pytestCheckHook,}:
buildPythonPackage (finalAttrs: { pname = "distrax"; version = "0.1.9"; pyproject = true; __structuredAttrs = true;
src = fetchFromGitHub { owner = "google-deepmind"; repo = "distrax"; tag = "v${finalAttrs.version}"; hash = "sha256-mX05qWyGTye+ZIXzU+W8ICz691UgVNIYXFN7oJHPssc="; };
build-system = [ flit-core ];
pythonRemoveDeps = [ "tfp-nightly" ]; dependencies = [ absl-py chex jax jaxlib numpy tensorflow-probability ];
nativeCheckInputs = [ dm-haiku mock pytest-xdist pytestCheckHook ];
pythonImportsCheck = [ "distrax" ];
disabledTests = [ # Flaky: AssertionError: 1 not less than 0.7000000000000001 "test_von_mises_sample_gradient" "test_von_mises_sample_moments" "test_von_mises_sample_uniform_ks_test"
# Flaky: AssertionError: Not equal to tolerance "StraightThroughTest" "test_composite_methods_are_consistent__with_jit" ];
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/hmm_test.py" "distrax/_src/utils/transformations_test.py" ];
meta = { description = "Probability distributions in JAX"; homepage = "https://github.com/deepmind/distrax"; changelog = "https://github.com/google-deepmind/distrax/releases/tag/${finalAttrs.src.tag}"; license = lib.licenses.asl20; maintainers = with lib.maintainers; [ onny ]; };})