pkgs/development/python-modules/ml-dtypes/default.nix

   1 { lib
   2 , buildPythonPackage
   3 , pythonOlder
   4 , fetchFromGitHub
   5 , setuptools
   6 , numpy
   7 , pytestCheckHook
   8 , absl-py
   9 }:
  10
  11 buildPythonPackage rec {
  12   pname = "ml-dtypes";
  13   version = "0.3.2";
  14   pyproject = true;
  15
  16   disabled = pythonOlder "3.9";
  17
  18   src = fetchFromGitHub {
  19     owner = "jax-ml";
  20     repo = "ml_dtypes";
  21     rev = "refs/tags/v${version}";
  22     hash = "sha256-epWunA5FULmCuTABl3uckFuNaSEpqJxtp0n0loCb6Q0=";
  23     # Since this upstream patch (https://github.com/jax-ml/ml_dtypes/commit/1bfd097e794413b0d465fa34f2eff0f3828ff521),
  24     # the attempts to use the nixpkgs packaged eigen dependency have failed.
  25     # Hence, we rely on the bundled eigen library.
  26     fetchSubmodules = true;
  27   };
  28
  29   postPatch = ''
  30     substituteInPlace pyproject.toml \
  31       --replace "numpy~=1.21.2" "numpy" \
  32       --replace "numpy~=1.23.3" "numpy" \
  33       --replace "numpy~=1.26.0" "numpy" \
  34       --replace "setuptools~=68.1.0" "setuptools"
  35   '';
  36
  37   nativeBuildInputs = [
  38     setuptools
  39   ];
  40
  41   propagatedBuildInputs = [
  42     numpy
  43   ];
  44
  45   nativeCheckInputs = [
  46     pytestCheckHook
  47     absl-py
  48   ];
  49
  50   preCheck = ''
  51     # remove src module, so tests use the installed module instead
  52     mv ./ml_dtypes/tests ./tests
  53     rm -rf ./ml_dtypes
  54   '';
  55
  56   pythonImportsCheck = [
  57     "ml_dtypes"
  58   ];
  59
  60   meta = with lib; {
  61     description = "A stand-alone implementation of several NumPy dtype extensions used in machine learning libraries";
  62     homepage = "https://github.com/jax-ml/ml_dtypes";
  63     changelog = "https://github.com/jax-ml/ml_dtypes/releases/tag/v${version}";
  64     license = licenses.asl20;
  65     maintainers = with maintainers; [ GaetanLepage samuela ];
  66   };
  67 }