pyrox.dev
1{
2 lib,
3 stdenv,
4 buildPythonPackage,
5 fetchFromGitHub,
6 rustPlatform,
7
8 # optional-dependencies
9 numpy,
10 torch,
11 tensorflow,
12 flax,
13 jax,
14 mlx,
15 paddlepaddle,
16 h5py,
17 huggingface-hub,
18 setuptools-rust,
19 pytest,
20 pytest-benchmark,
21 hypothesis,
22 fsspec,
23
24 # tests
25 pytestCheckHook,
26}:
27
28buildPythonPackage rec {
29 pname = "safetensors";
30 version = "0.6.2";
31 pyproject = true;
32
33 src = fetchFromGitHub {
34 owner = "huggingface";
35 repo = "safetensors";
36 tag = "v${version}";
37 hash = "sha256-IyKk29jMAbYW+16mrpqQWjnsmNFEvUwkB048AAx/Cvw=";
38 };
39
40 sourceRoot = "${src.name}/bindings/python";
41
42 cargoDeps = rustPlatform.fetchCargoVendor {
43 inherit
44 pname
45 version
46 src
47 sourceRoot
48 ;
49 hash = "sha256-+92fCILZwk/TknGXgR9lRN55WnmkgUJfCszFthstzXs=";
50 };
51
52 nativeBuildInputs = [
53 rustPlatform.cargoSetupHook
54 rustPlatform.maturinBuildHook
55 ];
56
57 optional-dependencies = lib.fix (self: {
58 numpy = [ numpy ];
59 torch = self.numpy ++ [
60 torch
61 ];
62 tensorflow = self.numpy ++ [
63 tensorflow
64 ];
65 pinned-tf = self.numpy ++ [
66 tensorflow
67 ];
68 jax = self.numpy ++ [
69 flax
70 jax
71 ];
72 mlx = [
73 mlx
74 ];
75 paddlepaddle = self.numpy ++ [
76 paddlepaddle
77 ];
78 testing = self.numpy ++ [
79 h5py
80 huggingface-hub
81 setuptools-rust
82 pytest
83 pytest-benchmark
84 hypothesis
85 fsspec
86 ];
87 all = self.torch ++ self.numpy ++ self.pinned-tf ++ self.jax ++ self.paddlepaddle ++ self.testing;
88 dev = self.all;
89 });
90
91 nativeCheckInputs = [
92 h5py
93 numpy
94 pytestCheckHook
95 torch
96 fsspec
97 ];
98
99 enabledTestPaths = [ "tests" ];
100
101 disabledTests = [
102 # AttributeError: module 'torch' has no attribute 'float4_e2m1fn_x2'
103 "test_odd_dtype_fp4"
104
105 # AssertionError: 'No such file or directory: notafile' != 'No such file or directory: "notafile"'
106 "test_file_not_found"
107
108 # AssertionError:
109 # 'Erro[41 chars] 5]: index 20 out of bounds for tensor dimension #1 of size 5'
110 # != 'Erro[41 chars] 5]: SliceOutOfRange { dim_index: 1, asked: 20, dim_size: 5 }'
111 "test_numpy_slice"
112 ];
113
114 # don't require PaddlePaddle (not in Nixpkgs), Flax, or Tensorflow (onerous) to run tests:
115 disabledTestPaths = [
116 "tests/test_flax_comparison.py"
117 "tests/test_paddle_comparison.py"
118 "tests/test_tf_comparison.py"
119 ]
120 ++ lib.optionals stdenv.hostPlatform.isDarwin [
121 # don't require mlx (not in Nixpkgs) to run tests
122 "tests/test_mlx_comparison.py"
123 ];
124
125 pythonImportsCheck = [ "safetensors" ];
126
127 meta = {
128 homepage = "https://github.com/huggingface/safetensors";
129 description = "Fast (zero-copy) and safe (unlike pickle) format for storing tensors";
130 changelog = "https://github.com/huggingface/safetensors/releases/tag/v${version}";
131 license = lib.licenses.asl20;
132 maintainers = with lib.maintainers; [ bcdarwin ];
133 };
134}