tjh.dev
nixpkgs mirror (for testing)
github.com/NixOS/nixpkgs
nix
1{
2 lib,
3 stdenv,
4 buildPythonPackage,
5 fetchFromGitHub,
6
7 # build-system
8 setuptools,
9 cython,
10 versioneer,
11
12 # dependencies
13 cons,
14 etuples,
15 filelock,
16 logical-unification,
17 minikanren,
18 numba,
19 numpy,
20 scipy,
21
22 # tests
23 jax,
24 jaxlib,
25 pytest-benchmark,
26 pytest-mock,
27 pytestCheckHook,
28 tensorflow-probability,
29 writableTmpDirAsHomeHook,
30
31 nix-update-script,
32}:
33
34buildPythonPackage (finalAttrs: {
35 pname = "pytensor";
36 version = "2.38.2";
37 pyproject = true;
38
39 src = fetchFromGitHub {
40 owner = "pymc-devs";
41 repo = "pytensor";
42 tag = "rel-${finalAttrs.version}";
43 postFetch = ''
44 sed -i 's/git_refnames = "[^"]*"/git_refnames = " (tag: ${finalAttrs.src.tag})"/' $out/pytensor/_version.py
45 '';
46 hash = "sha256-BKyaApIijxuJ0gNNXqahDOMW3rpF6+qgoCEpWj6Uz5g=";
47 };
48
49 build-system = [
50 setuptools
51 cython
52 versioneer
53 ];
54
55 dependencies = [
56 cons
57 etuples
58 filelock
59 logical-unification
60 minikanren
61 numba
62 numpy
63 scipy
64 setuptools
65 ];
66
67 nativeCheckInputs = [
68 jax
69 jaxlib
70 numba
71 pytest-benchmark
72 pytest-mock
73 pytestCheckHook
74 tensorflow-probability
75 writableTmpDirAsHomeHook
76 ];
77
78 pytestFlags = [ "--benchmark-disable" ];
79
80 pythonImportsCheck = [ "pytensor" ];
81
82 # Ensure that the installed package is used instead of the source files from the current workdir
83 preCheck = ''
84 rm -rf pytensor
85 '';
86
87 disabledTests = [
88 # TypeError: jax_funcified_fgraph() takes 2 positional arguments but 3 were given
89 "test_jax_Reshape_shape_graph_input"
90
91 # AssertionError: equal_computations failed
92 "test_infer_shape_db_handles_xtensor_lowering"
93 ]
94 ++ lib.optionals stdenv.hostPlatform.isDarwin [
95 # Numerical assertion error
96 # tests.unittest_tools.WrongValue: WrongValue
97 "test_op_sd"
98 "test_op_ss"
99
100 # AssertionError: equal_computations failed
101 "test_infer_shape_db_handles_xtensor_lowering"
102
103 # pytensor.link.c.exceptions.CompileError: Compilation failed (return status=1)
104 "OpFromGraph"
105 "add"
106 "cls_ofg1"
107 "direct"
108 "multiply"
109 "test_AddDS"
110 "test_AddSD"
111 "test_AddSS"
112 "test_MulDS"
113 "test_MulSD"
114 "test_MulSS"
115 "test_NoOutputFromInplace"
116 "test_OpFromGraph"
117 "test_adv_sub1_sparse_grad"
118 "test_alloc"
119 "test_binary"
120 "test_borrow_input"
121 "test_borrow_output"
122 "test_cache_race_condition"
123 "test_check_for_aliased_inputs"
124 "test_clinker_literal_cache"
125 "test_csm_grad"
126 "test_csm_unsorted"
127 "test_csr_dense_grad"
128 "test_debugprint"
129 "test_ellipsis_einsum"
130 "test_empty_elemwise"
131 "test_flatten"
132 "test_fprop"
133 "test_get_item_list_grad"
134 "test_grad"
135 "test_infer_shape"
136 "test_jax_pad"
137 "test_kron"
138 "test_masked_input"
139 "test_max"
140 "test_modes"
141 "test_mul_s_v_grad"
142 "test_multiple_outputs"
143 "test_nnet"
144 "test_not_inplace"
145 "test_numba_Cholesky_grad"
146 "test_numba_pad"
147 "test_optimizations_preserved"
148 "test_overided_function"
149 "test_potential_output_aliasing_induced_by_updates"
150 "test_profiling"
151 "test_rebuild_strict"
152 "test_runtime_broadcast_c"
153 "test_scan_err1"
154 "test_scan_err2"
155 "test_shared"
156 "test_size_implied_by_broadcasted_parameters"
157 "test_solve_triangular_grad"
158 "test_structured_add_s_v_grad"
159 "test_structureddot_csc_grad"
160 "test_structureddot_csr_grad"
161 "test_sum"
162 "test_swap_SharedVariable_with_given"
163 "test_test_value_op"
164 "test_unary"
165 "test_unbroadcast"
166 "test_update_equiv"
167 "test_update_same"
168 ];
169
170 disabledTestPaths = [
171 # Don't run the most compute-intense tests
172 "tests/scan/"
173 "tests/tensor/"
174 ];
175
176 passthru.updateScript = nix-update-script {
177 extraArgs = [
178 "--version-regex"
179 "rel-(.+)"
180 ];
181 };
182
183 meta = {
184 description = "Python library to define, optimize, and efficiently evaluate mathematical expressions involving multi-dimensional arrays";
185 mainProgram = "pytensor-cache";
186 homepage = "https://github.com/pymc-devs/pytensor";
187 changelog = "https://github.com/pymc-devs/pytensor/releases/tag/${finalAttrs.src.tag}";
188 license = lib.licenses.bsd3;
189 maintainers = with lib.maintainers; [
190 bcdarwin
191 ];
192 };
193})