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