pkgs/development/python-modules/pytensor/default.nix

   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   tensorflow-probability,
  26
  27   nix-update-script,
  28 }:
  29
  30 buildPythonPackage rec {
  31   pname = "pytensor";
  32   version = "2.25.5";
  33   pyproject = true;
  34
  35   src = fetchFromGitHub {
  36     owner = "pymc-devs";
  37     repo = "pytensor";
  38     rev = "refs/tags/rel-${version}";
  39     hash = "sha256-tEIDuMAY+rMBqzHLnhbSqOoWzjIbZKCWWMx/uej2pOU=";
  40   };
  41
  42   pythonRelaxDeps = [
  43     "scipy"
  44   ];
  45
  46   build-system = [
  47     cython
  48     versioneer
  49   ];
  50
  51   dependencies = [
  52     cons
  53     etuples
  54     filelock
  55     logical-unification
  56     minikanren
  57     numpy
  58     scipy
  59   ];
  60
  61   nativeCheckInputs = [
  62     jax
  63     jaxlib
  64     numba
  65     pytest-mock
  66     pytestCheckHook
  67     tensorflow-probability
  68   ];
  69
  70   preBuild = ''
  71     export HOME=$(mktemp -d)
  72   '';
  73
  74   pythonImportsCheck = [ "pytensor" ];
  75
  76   disabledTests = [
  77     # benchmarks (require pytest-benchmark):
  78     "test_elemwise_speed"
  79     "test_fused_elemwise_benchmark"
  80     "test_logsumexp_benchmark"
  81     "test_scan_multiple_output"
  82     "test_vector_taps_benchmark"
  83   ];
  84
  85   disabledTestPaths = [
  86     # Don't run the most compute-intense tests
  87     "tests/scan/"
  88     "tests/tensor/"
  89     "tests/sparse/sandbox/"
  90   ];
  91
  92   passthru.updateScript = nix-update-script {
  93     extraArgs = [
  94       "--version-regex"
  95       "rel-(.+)"
  96     ];
  97   };
  98
  99   meta = {
 100     description = "Python library to define, optimize, and efficiently evaluate mathematical expressions involving multi-dimensional arrays";
 101     mainProgram = "pytensor-cache";
 102     homepage = "https://github.com/pymc-devs/pytensor";
 103     changelog = "https://github.com/pymc-devs/pytensor/releases";
 104     license = lib.licenses.bsd3;
 105     maintainers = with lib.maintainers; [
 106       bcdarwin
 107       ferrine
 108     ];
 109   };
 110 }