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

   1 { lib
   2 , buildPythonPackage
   3 , fetchFromGitHub
   4 , hatchling
   5 , numpy
   6 , typeguard
   7 , typing-extensions
   8 , cloudpickle
   9 , equinox
  10 , jax
  11 , jaxlib
  12 , torch
  13 , pytestCheckHook
  14 }:
  15
  16 let
  17   self = buildPythonPackage rec {
  18     pname = "jaxtyping";
  19     version = "0.2.23";
  20     pyproject = true;
  21
  22     src = fetchFromGitHub {
  23       owner = "google";
  24       repo = "jaxtyping";
  25       rev = "refs/tags/v${version}";
  26       hash = "sha256-22dIuIjFgqRmV9AQok02skVt7fm17/WpzBm3FrJ6/zs=";
  27     };
  28
  29     nativeBuildInputs = [
  30       hatchling
  31     ];
  32
  33     propagatedBuildInputs = [
  34       numpy
  35       typeguard
  36       typing-extensions
  37     ];
  38
  39     nativeCheckInputs = [
  40       cloudpickle
  41       equinox
  42       jax
  43       jaxlib
  44       pytestCheckHook
  45       torch
  46     ];
  47
  48     doCheck = false;
  49
  50     # Enable tests via passthru to avoid cyclic dependency with equinox.
  51     passthru.tests = {
  52       check = self.overridePythonAttrs { doCheck = true; };
  53     };
  54
  55     pythonImportsCheck = [ "jaxtyping" ];
  56
  57     meta = with lib; {
  58       description = "Type annotations and runtime checking for JAX arrays and PyTrees";
  59       homepage = "https://github.com/google/jaxtyping";
  60       license = licenses.mit;
  61       maintainers = with maintainers; [ GaetanLepage ];
  62     };
  63   };
  64  in self