diff --git a/pkgs/development/python-modules/safetensors/default.nix b/pkgs/development/python-modules/safetensors/default.nix index eb7d1b502d91..e82879ccf563 100644 --- a/pkgs/development/python-modules/safetensors/default.nix +++ b/pkgs/development/python-modules/safetensors/default.nix @@ -5,11 +5,23 @@ fetchFromGitHub, rustPlatform, - # tests - h5py, + # optional-dependencies numpy, - pytestCheckHook, torch, + tensorflow, + flax, + jax, + mlx, + paddlepaddle, + h5py, + huggingface-hub, + setuptools-rust, + pytest, + pytest-benchmark, + hypothesis, + + # tests + pytestCheckHook, }: buildPythonPackage rec { @@ -36,6 +48,39 @@ buildPythonPackage rec { rustPlatform.maturinBuildHook ]; + optional-dependencies = lib.fix (self: { + numpy = [ numpy ]; + torch = self.numpy ++ [ + torch + ]; + tensorflow = self.numpy ++ [ + tensorflow + ]; + pinned-tf = self.numpy ++ [ + tensorflow + ]; + jax = self.numpy ++ [ + flax + jax + ]; + mlx = [ + mlx + ]; + paddlepaddle = self.numpy ++ [ + paddlepaddle + ]; + testing = self.numpy ++ [ + h5py + huggingface-hub + setuptools-rust + pytest + pytest-benchmark + hypothesis + ]; + all = self.torch ++ self.numpy ++ self.pinned-tf ++ self.jax ++ self.paddlepaddle ++ self.testing; + dev = self.all; + }); + nativeCheckInputs = [ h5py numpy