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

   1 { lib
   2 , buildPythonPackage
   3 , pythonOlder
   4 , fetchFromGitHub
   5 , hatchling
   6 , pythonRelaxDepsHook
   7 , numpy
   8 , typeguard
   9 , typing-extensions
  10 , cloudpickle
  11 , equinox
  12 , ipython
  13 , jax
  14 , jaxlib
  15 , pytestCheckHook
  16 , tensorflow
  17 , torch
  18 }:
  19
  20 let
  21   self = buildPythonPackage rec {
  22     pname = "jaxtyping";
  23     version = "0.2.28";
  24     pyproject = true;
  25
  26     disabled = pythonOlder "3.9";
  27
  28     src = fetchFromGitHub {
  29       owner = "google";
  30       repo = "jaxtyping";
  31       rev = "refs/tags/v${version}";
  32       hash = "sha256-xDFrgPecUIfCACg/xkMQ8G1+6hNiUUDg9eCZKNpNfzs=";
  33     };
  34
  35     nativeBuildInputs = [
  36       hatchling
  37       pythonRelaxDepsHook
  38     ];
  39
  40     propagatedBuildInputs = [
  41       numpy
  42       typeguard
  43       typing-extensions
  44     ];
  45
  46     pythonRelaxDeps = [
  47       "typeguard"
  48     ];
  49
  50     nativeCheckInputs = [
  51       cloudpickle
  52       equinox
  53       ipython
  54       jax
  55       jaxlib
  56       pytestCheckHook
  57       tensorflow
  58       torch
  59     ];
  60
  61     doCheck = false;
  62
  63     # Enable tests via passthru to avoid cyclic dependency with equinox.
  64     passthru.tests = {
  65       check = self.overridePythonAttrs {
  66         # We disable tests because they complain about the version of typeguard being too new.
  67         doCheck = false;
  68         catchConflicts = false;
  69       };
  70     };
  71
  72     pythonImportsCheck = [ "jaxtyping" ];
  73
  74     meta = with lib; {
  75       description = "Type annotations and runtime checking for JAX arrays and PyTrees";
  76       homepage = "https://github.com/google/jaxtyping";
  77       license = licenses.mit;
  78       maintainers = with maintainers; [ GaetanLepage ];
  79     };
  80   };
  81  in self