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

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