pyrox.dev
1{
2 lib,
3 buildPythonPackage,
4 fetchFromGitHub,
5
6 # build-system
7 flit-core,
8
9 # dependencies
10 absl-py,
11 chex,
12 jax,
13 jaxlib,
14 numpy,
15 tensorflow-probability,
16
17 # tests
18 dm-haiku,
19 pytest-xdist,
20 pytestCheckHook,
21}:
22
23buildPythonPackage rec {
24 pname = "distrax";
25 version = "0.1.7";
26 pyproject = true;
27
28 src = fetchFromGitHub {
29 owner = "google-deepmind";
30 repo = "distrax";
31 tag = "v${version}";
32 hash = "sha256-R6rGGNzup3O6eZ2z4vygYWTjroE/Irt3aog8Op+0hco=";
33 };
34
35 build-system = [
36 flit-core
37 ];
38
39 dependencies = [
40 absl-py
41 chex
42 jax
43 jaxlib
44 numpy
45 tensorflow-probability
46 ];
47
48 nativeCheckInputs = [
49 dm-haiku
50 pytest-xdist
51 pytestCheckHook
52 ];
53
54 pythonImportsCheck = [ "distrax" ];
55
56 disabledTests = [
57 # Flaky: AssertionError: 1 not less than 0.7000000000000001
58 "test_von_mises_sample_uniform_ks_test"
59
60 # Flaky: AssertionError: Not equal to tolerance
61 "test_composite_methods_are_consistent__with_jit"
62
63 # NotImplementedError: Primitive 'square' does not have a registered inverse.
64 "test_against_tfp_bijectors_square"
65 "test_log_dets_square__with_device"
66 "test_log_dets_square__without_device"
67 "test_log_dets_square__without_jit"
68
69 # AssertionError on numerical values
70 # Reported upstream in https://github.com/google-deepmind/distrax/issues/267
71 "test_method_with_input_unnormalized_probs__with_device"
72 "test_method_with_input_unnormalized_probs__with_jit"
73 "test_method_with_input_unnormalized_probs__without_device"
74 "test_method_with_input_unnormalized_probs__without_jit"
75 "test_method_with_value_1d"
76 "test_nested_distributions__with_device"
77 "test_nested_distributions__without_device"
78 "test_nested_distributions__with_jit"
79 "test_nested_distributions__without_jit"
80 "test_stability__with_device"
81 "test_stability__with_jit"
82 "test_stability__without_device"
83 "test_stability__without_jit"
84 "test_von_mises_sample_gradient"
85 "test_von_mises_sample_moments"
86 ];
87
88 disabledTestPaths = [
89 # Since jax 0.6.0:
90 # TypeError: <lambda>() got an unexpected keyword argument 'accuracy'
91 "distrax/_src/bijectors/lambda_bijector_test.py"
92
93 # TypeErrors
94 "distrax/_src/bijectors/tfp_compatible_bijector_test.py"
95 "distrax/_src/distributions/distribution_from_tfp_test.py"
96 "distrax/_src/distributions/laplace_test.py"
97 "distrax/_src/distributions/multinomial_test.py"
98 "distrax/_src/distributions/mvn_diag_plus_low_rank_test.py"
99 "distrax/_src/distributions/mvn_kl_test.py"
100 "distrax/_src/distributions/straight_through_test.py"
101 "distrax/_src/distributions/tfp_compatible_distribution_test.py"
102 "distrax/_src/distributions/transformed_test.py"
103 "distrax/_src/distributions/uniform_test.py"
104 "distrax/_src/utils/transformations_test.py"
105 # https://github.com/google-deepmind/distrax/pull/270
106 "distrax/_src/distributions/deterministic_test.py"
107 "distrax/_src/distributions/epsilon_greedy_test.py"
108 "distrax/_src/distributions/gamma_test.py"
109 "distrax/_src/distributions/greedy_test.py"
110 "distrax/_src/distributions/gumbel_test.py"
111 "distrax/_src/distributions/logistic_test.py"
112 "distrax/_src/distributions/log_stddev_normal_test.py"
113 "distrax/_src/distributions/mvn_diag_test.py"
114 "distrax/_src/distributions/mvn_full_covariance_test.py"
115 "distrax/_src/distributions/mvn_tri_test.py"
116 "distrax/_src/distributions/one_hot_categorical_test.py"
117 "distrax/_src/distributions/softmax_test.py"
118 "distrax/_src/utils/hmm_test.py"
119 ];
120
121 meta = {
122 description = "Probability distributions in JAX";
123 homepage = "https://github.com/deepmind/distrax";
124 changelog = "https://github.com/google-deepmind/distrax/releases/tag/v${version}";
125 license = lib.licenses.asl20;
126 maintainers = with lib.maintainers; [ onny ];
127 badPlatforms = [
128 # SystemError: nanobind::detail::nb_func_error_except(): exception could not be translated!
129 lib.systems.inspect.patterns.isDarwin
130 ];
131 };
132}