tjh.dev
1{ stdenv
2, lib
3, buildPythonPackage
4, cons
5, cython
6, etuples
7, fetchFromGitHub
8, filelock
9, jax
10, jaxlib
11, logical-unification
12, minikanren
13, numba
14, numba-scipy
15, numpy
16, pytestCheckHook
17, pythonOlder
18, scipy
19, typing-extensions
20}:
21
22buildPythonPackage rec {
23 pname = "pytensor";
24 version = "2.11.3";
25 format = "setuptools";
26
27 disabled = pythonOlder "3.7";
28
29 src = fetchFromGitHub {
30 owner = "pymc-devs";
31 repo = pname;
32 rev = "refs/tags/rel-${version}";
33 hash = "sha256-4GDur8S19i8pZkywKHZUelmd2e0jZmC5HzF7o2esDl4=";
34 };
35
36 nativeBuildInputs = [
37 cython
38 ];
39
40 propagatedBuildInputs = [
41 cons
42 etuples
43 filelock
44 logical-unification
45 minikanren
46 numpy
47 scipy
48 typing-extensions
49 ];
50
51 checkInputs = [
52 jax
53 jaxlib
54 numba
55 numba-scipy
56 pytestCheckHook
57 ];
58
59 postPatch = ''
60 substituteInPlace setup.cfg \
61 --replace "--durations=50" ""
62 '';
63
64 preBuild = ''
65 export HOME=$(mktemp -d)
66 '';
67
68 pythonImportsCheck = [
69 "pytensor"
70 ];
71
72 disabledTests = [
73 # benchmarks (require pytest-benchmark):
74 "test_elemwise_speed"
75 "test_fused_elemwise_benchmark"
76 "test_logsumexp_benchmark"
77 "test_scan_multiple_output"
78 "test_vector_taps_benchmark"
79 ];
80
81 disabledTestPaths = [
82 # Don't run the most compute-intense tests
83 "tests/scan/"
84 "tests/tensor/"
85 "tests/sandbox/"
86 "tests/sparse/sandbox/"
87 ];
88
89 meta = with lib; {
90 description = "Python library to define, optimize, and efficiently evaluate mathematical expressions involving multi-dimensional arrays";
91 homepage = "https://github.com/pymc-devs/pytensor";
92 changelog = "https://github.com/pymc-devs/pytensor/releases";
93 license = licenses.bsd3;
94 maintainers = with maintainers; [ bcdarwin ];
95 broken = (stdenv.isLinux && stdenv.isAarch64);
96 };
97}