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

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