safetensors 0.5.2__tar.gz → 0.5.3__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of safetensors might be problematic. Click here for more details.

Files changed (52) hide show
  1. {safetensors-0.5.2 → safetensors-0.5.3}/PKG-INFO +20 -20
  2. {safetensors-0.5.2 → safetensors-0.5.3}/bindings/python/Cargo.lock +33 -33
  3. {safetensors-0.5.2 → safetensors-0.5.3}/bindings/python/Cargo.toml +2 -1
  4. {safetensors-0.5.2 → safetensors-0.5.3}/bindings/python/src/lib.rs +11 -3
  5. {safetensors-0.5.2 → safetensors-0.5.3}/bindings/python/tests/test_pt_comparison.py +18 -0
  6. {safetensors-0.5.2 → safetensors-0.5.3}/bindings/python/tests/test_simple.py +15 -1
  7. {safetensors-0.5.2 → safetensors-0.5.3}/safetensors/Cargo.toml +2 -1
  8. {safetensors-0.5.2 → safetensors-0.5.3}/safetensors/src/slice.rs +59 -2
  9. {safetensors-0.5.2 → safetensors-0.5.3}/safetensors/src/tensor.rs +4 -0
  10. {safetensors-0.5.2 → safetensors-0.5.3}/bindings/python/.gitignore +0 -0
  11. {safetensors-0.5.2 → safetensors-0.5.3}/bindings/python/MANIFEST.in +0 -0
  12. {safetensors-0.5.2 → safetensors-0.5.3}/bindings/python/Makefile +0 -0
  13. {safetensors-0.5.2 → safetensors-0.5.3}/bindings/python/README.md +0 -0
  14. {safetensors-0.5.2 → safetensors-0.5.3}/bindings/python/benches/test_flax.py +0 -0
  15. {safetensors-0.5.2 → safetensors-0.5.3}/bindings/python/benches/test_mlx.py +0 -0
  16. {safetensors-0.5.2 → safetensors-0.5.3}/bindings/python/benches/test_paddle.py +0 -0
  17. {safetensors-0.5.2 → safetensors-0.5.3}/bindings/python/benches/test_pt.py +0 -0
  18. {safetensors-0.5.2 → safetensors-0.5.3}/bindings/python/benches/test_tf.py +0 -0
  19. {safetensors-0.5.2 → safetensors-0.5.3}/bindings/python/convert.py +0 -0
  20. {safetensors-0.5.2 → safetensors-0.5.3}/bindings/python/convert_all.py +0 -0
  21. {safetensors-0.5.2 → safetensors-0.5.3}/bindings/python/fuzz.py +0 -0
  22. {safetensors-0.5.2 → safetensors-0.5.3}/bindings/python/py_src/safetensors/__init__.py +0 -0
  23. {safetensors-0.5.2 → safetensors-0.5.3}/bindings/python/py_src/safetensors/__init__.pyi +0 -0
  24. {safetensors-0.5.2 → safetensors-0.5.3}/bindings/python/py_src/safetensors/flax.py +0 -0
  25. {safetensors-0.5.2 → safetensors-0.5.3}/bindings/python/py_src/safetensors/mlx.py +0 -0
  26. {safetensors-0.5.2 → safetensors-0.5.3}/bindings/python/py_src/safetensors/numpy.py +0 -0
  27. {safetensors-0.5.2 → safetensors-0.5.3}/bindings/python/py_src/safetensors/paddle.py +0 -0
  28. {safetensors-0.5.2 → safetensors-0.5.3}/bindings/python/py_src/safetensors/py.typed +0 -0
  29. {safetensors-0.5.2 → safetensors-0.5.3}/bindings/python/py_src/safetensors/tensorflow.py +0 -0
  30. {safetensors-0.5.2 → safetensors-0.5.3}/bindings/python/py_src/safetensors/torch.py +0 -0
  31. {safetensors-0.5.2 → safetensors-0.5.3}/bindings/python/setup.cfg +0 -0
  32. {safetensors-0.5.2 → safetensors-0.5.3}/bindings/python/stub.py +0 -0
  33. {safetensors-0.5.2 → safetensors-0.5.3}/bindings/python/tests/data/__init__.py +0 -0
  34. {safetensors-0.5.2 → safetensors-0.5.3}/bindings/python/tests/test_flax_comparison.py +0 -0
  35. {safetensors-0.5.2 → safetensors-0.5.3}/bindings/python/tests/test_mlx_comparison.py +0 -0
  36. {safetensors-0.5.2 → safetensors-0.5.3}/bindings/python/tests/test_paddle_comparison.py +0 -0
  37. {safetensors-0.5.2 → safetensors-0.5.3}/bindings/python/tests/test_pt_model.py +0 -0
  38. {safetensors-0.5.2 → safetensors-0.5.3}/bindings/python/tests/test_tf_comparison.py +0 -0
  39. {safetensors-0.5.2 → safetensors-0.5.3}/py_src/safetensors/__init__.py +0 -0
  40. {safetensors-0.5.2 → safetensors-0.5.3}/py_src/safetensors/__init__.pyi +0 -0
  41. {safetensors-0.5.2 → safetensors-0.5.3}/py_src/safetensors/flax.py +0 -0
  42. {safetensors-0.5.2 → safetensors-0.5.3}/py_src/safetensors/mlx.py +0 -0
  43. {safetensors-0.5.2 → safetensors-0.5.3}/py_src/safetensors/numpy.py +0 -0
  44. {safetensors-0.5.2 → safetensors-0.5.3}/py_src/safetensors/paddle.py +0 -0
  45. {safetensors-0.5.2 → safetensors-0.5.3}/py_src/safetensors/py.typed +0 -0
  46. {safetensors-0.5.2 → safetensors-0.5.3}/py_src/safetensors/tensorflow.py +0 -0
  47. {safetensors-0.5.2 → safetensors-0.5.3}/py_src/safetensors/torch.py +0 -0
  48. {safetensors-0.5.2 → safetensors-0.5.3}/pyproject.toml +0 -0
  49. {safetensors-0.5.2 → safetensors-0.5.3}/safetensors/LICENSE +0 -0
  50. {safetensors-0.5.2 → safetensors-0.5.3}/safetensors/README.md +0 -0
  51. {safetensors-0.5.2 → safetensors-0.5.3}/safetensors/benches/benchmark.rs +0 -0
  52. {safetensors-0.5.2 → safetensors-0.5.3}/safetensors/src/lib.rs +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: safetensors
3
- Version: 0.5.2
3
+ Version: 0.5.3
4
4
  Classifier: Development Status :: 5 - Production/Stable
5
5
  Classifier: Intended Audience :: Developers
6
6
  Classifier: Intended Audience :: Education
@@ -14,31 +14,31 @@ Classifier: Programming Language :: Python :: 3.9
14
14
  Classifier: Programming Language :: Python :: 3.10
15
15
  Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
16
16
  Classifier: Typing :: Typed
17
- Requires-Dist: numpy >=1.21.6 ; extra == 'numpy'
17
+ Requires-Dist: numpy>=1.21.6 ; extra == 'numpy'
18
18
  Requires-Dist: safetensors[numpy] ; extra == 'torch'
19
- Requires-Dist: torch >=1.10 ; extra == 'torch'
19
+ Requires-Dist: torch>=1.10 ; extra == 'torch'
20
20
  Requires-Dist: safetensors[numpy] ; extra == 'tensorflow'
21
- Requires-Dist: tensorflow >=2.11.0 ; extra == 'tensorflow'
21
+ Requires-Dist: tensorflow>=2.11.0 ; extra == 'tensorflow'
22
22
  Requires-Dist: safetensors[numpy] ; extra == 'pinned-tf'
23
- Requires-Dist: tensorflow ==2.18.0 ; extra == 'pinned-tf'
23
+ Requires-Dist: tensorflow==2.18.0 ; extra == 'pinned-tf'
24
24
  Requires-Dist: safetensors[numpy] ; extra == 'jax'
25
- Requires-Dist: flax >=0.6.3 ; extra == 'jax'
26
- Requires-Dist: jax >=0.3.25 ; extra == 'jax'
27
- Requires-Dist: jaxlib >=0.3.25 ; extra == 'jax'
28
- Requires-Dist: mlx >=0.0.9 ; extra == 'mlx'
25
+ Requires-Dist: flax>=0.6.3 ; extra == 'jax'
26
+ Requires-Dist: jax>=0.3.25 ; extra == 'jax'
27
+ Requires-Dist: jaxlib>=0.3.25 ; extra == 'jax'
28
+ Requires-Dist: mlx>=0.0.9 ; extra == 'mlx'
29
29
  Requires-Dist: safetensors[numpy] ; extra == 'paddlepaddle'
30
- Requires-Dist: paddlepaddle >=2.4.1 ; extra == 'paddlepaddle'
31
- Requires-Dist: black ==22.3 ; extra == 'quality'
32
- Requires-Dist: click ==8.0.4 ; extra == 'quality'
33
- Requires-Dist: isort >=5.5.4 ; extra == 'quality'
34
- Requires-Dist: flake8 >=3.8.3 ; extra == 'quality'
30
+ Requires-Dist: paddlepaddle>=2.4.1 ; extra == 'paddlepaddle'
31
+ Requires-Dist: black==22.3 ; extra == 'quality'
32
+ Requires-Dist: click==8.0.4 ; extra == 'quality'
33
+ Requires-Dist: isort>=5.5.4 ; extra == 'quality'
34
+ Requires-Dist: flake8>=3.8.3 ; extra == 'quality'
35
35
  Requires-Dist: safetensors[numpy] ; extra == 'testing'
36
- Requires-Dist: h5py >=3.7.0 ; extra == 'testing'
37
- Requires-Dist: huggingface-hub >=0.12.1 ; extra == 'testing'
38
- Requires-Dist: setuptools-rust >=1.5.2 ; extra == 'testing'
39
- Requires-Dist: pytest >=7.2.0 ; extra == 'testing'
40
- Requires-Dist: pytest-benchmark >=4.0.0 ; extra == 'testing'
41
- Requires-Dist: hypothesis >=6.70.2 ; extra == 'testing'
36
+ Requires-Dist: h5py>=3.7.0 ; extra == 'testing'
37
+ Requires-Dist: huggingface-hub>=0.12.1 ; extra == 'testing'
38
+ Requires-Dist: setuptools-rust>=1.5.2 ; extra == 'testing'
39
+ Requires-Dist: pytest>=7.2.0 ; extra == 'testing'
40
+ Requires-Dist: pytest-benchmark>=4.0.0 ; extra == 'testing'
41
+ Requires-Dist: hypothesis>=6.70.2 ; extra == 'testing'
42
42
  Requires-Dist: safetensors[torch] ; extra == 'all'
43
43
  Requires-Dist: safetensors[numpy] ; extra == 'all'
44
44
  Requires-Dist: safetensors[pinned-tf] ; extra == 'all'
@@ -1,6 +1,6 @@
1
1
  # This file is automatically @generated by Cargo.
2
2
  # It is not intended for manual editing.
3
- version = 4
3
+ version = 3
4
4
 
5
5
  [[package]]
6
6
  name = "autocfg"
@@ -34,9 +34,9 @@ checksum = "d75a2a4b1b190afb6f5425f10f6a8f959d2ea0b9c2b1d79553551850539e4674"
34
34
 
35
35
  [[package]]
36
36
  name = "libc"
37
- version = "0.2.169"
37
+ version = "0.2.170"
38
38
  source = "registry+https://github.com/rust-lang/crates.io-index"
39
- checksum = "b5aba8db14291edd000dfcc4d620c7ebfb122c613afb886ca8803fa4e128a20a"
39
+ checksum = "875b3680cb2f8f71bdcf9a30f38d48282f5d3c95cbf9b3fa57269bb5d5c06828"
40
40
 
41
41
  [[package]]
42
42
  name = "memchr"
@@ -64,30 +64,30 @@ dependencies = [
64
64
 
65
65
  [[package]]
66
66
  name = "once_cell"
67
- version = "1.20.2"
67
+ version = "1.20.3"
68
68
  source = "registry+https://github.com/rust-lang/crates.io-index"
69
- checksum = "1261fe7e33c73b354eab43b1273a57c8f967d0391e80353e51f764ac02cf6775"
69
+ checksum = "945462a4b81e43c4e3ba96bd7b49d834c6f61198356aa858733bc4acf3cbe62e"
70
70
 
71
71
  [[package]]
72
72
  name = "portable-atomic"
73
- version = "1.10.0"
73
+ version = "1.11.0"
74
74
  source = "registry+https://github.com/rust-lang/crates.io-index"
75
- checksum = "280dc24453071f1b63954171985a0b0d30058d287960968b9b2aca264c8d4ee6"
75
+ checksum = "350e9b48cbc6b0e028b0473b114454c6316e57336ee184ceab6e53f72c178b3e"
76
76
 
77
77
  [[package]]
78
78
  name = "proc-macro2"
79
- version = "1.0.92"
79
+ version = "1.0.93"
80
80
  source = "registry+https://github.com/rust-lang/crates.io-index"
81
- checksum = "37d3544b3f2748c54e147655edb5025752e2303145b5aefb3c3ea2c78b973bb0"
81
+ checksum = "60946a68e5f9d28b0dc1c21bb8a97ee7d018a8b322fa57838ba31cc878e22d99"
82
82
  dependencies = [
83
83
  "unicode-ident",
84
84
  ]
85
85
 
86
86
  [[package]]
87
87
  name = "pyo3"
88
- version = "0.23.3"
88
+ version = "0.23.5"
89
89
  source = "registry+https://github.com/rust-lang/crates.io-index"
90
- checksum = "e484fd2c8b4cb67ab05a318f1fd6fa8f199fcc30819f08f07d200809dba26c15"
90
+ checksum = "7778bffd85cf38175ac1f545509665d0b9b92a198ca7941f131f85f7a4f9a872"
91
91
  dependencies = [
92
92
  "cfg-if",
93
93
  "indoc",
@@ -103,9 +103,9 @@ dependencies = [
103
103
 
104
104
  [[package]]
105
105
  name = "pyo3-build-config"
106
- version = "0.23.3"
106
+ version = "0.23.5"
107
107
  source = "registry+https://github.com/rust-lang/crates.io-index"
108
- checksum = "dc0e0469a84f208e20044b98965e1561028180219e35352a2afaf2b942beff3b"
108
+ checksum = "94f6cbe86ef3bf18998d9df6e0f3fc1050a8c5efa409bf712e661a4366e010fb"
109
109
  dependencies = [
110
110
  "once_cell",
111
111
  "target-lexicon",
@@ -113,9 +113,9 @@ dependencies = [
113
113
 
114
114
  [[package]]
115
115
  name = "pyo3-ffi"
116
- version = "0.23.3"
116
+ version = "0.23.5"
117
117
  source = "registry+https://github.com/rust-lang/crates.io-index"
118
- checksum = "eb1547a7f9966f6f1a0f0227564a9945fe36b90da5a93b3933fc3dc03fae372d"
118
+ checksum = "e9f1b4c431c0bb1c8fb0a338709859eed0d030ff6daa34368d3b152a63dfdd8d"
119
119
  dependencies = [
120
120
  "libc",
121
121
  "pyo3-build-config",
@@ -123,9 +123,9 @@ dependencies = [
123
123
 
124
124
  [[package]]
125
125
  name = "pyo3-macros"
126
- version = "0.23.3"
126
+ version = "0.23.5"
127
127
  source = "registry+https://github.com/rust-lang/crates.io-index"
128
- checksum = "fdb6da8ec6fa5cedd1626c886fc8749bdcbb09424a86461eb8cdf096b7c33257"
128
+ checksum = "fbc2201328f63c4710f68abdf653c89d8dbc2858b88c5d88b0ff38a75288a9da"
129
129
  dependencies = [
130
130
  "proc-macro2",
131
131
  "pyo3-macros-backend",
@@ -135,9 +135,9 @@ dependencies = [
135
135
 
136
136
  [[package]]
137
137
  name = "pyo3-macros-backend"
138
- version = "0.23.3"
138
+ version = "0.23.5"
139
139
  source = "registry+https://github.com/rust-lang/crates.io-index"
140
- checksum = "38a385202ff5a92791168b1136afae5059d3ac118457bb7bc304c197c2d33e7d"
140
+ checksum = "fca6726ad0f3da9c9de093d6f116a93c1a38e417ed73bf138472cf4064f72028"
141
141
  dependencies = [
142
142
  "heck",
143
143
  "proc-macro2",
@@ -157,13 +157,13 @@ dependencies = [
157
157
 
158
158
  [[package]]
159
159
  name = "ryu"
160
- version = "1.0.18"
160
+ version = "1.0.19"
161
161
  source = "registry+https://github.com/rust-lang/crates.io-index"
162
- checksum = "f3cb5ba0dc43242ce17de99c180e96db90b235b8a9fdc9543c96d2209116bd9f"
162
+ checksum = "6ea1a2d0a644769cc99faa24c3ad26b379b786fe7c36fd3c546254801650e6dd"
163
163
 
164
164
  [[package]]
165
165
  name = "safetensors"
166
- version = "0.5.2"
166
+ version = "0.5.3"
167
167
  dependencies = [
168
168
  "serde",
169
169
  "serde_json",
@@ -171,7 +171,7 @@ dependencies = [
171
171
 
172
172
  [[package]]
173
173
  name = "safetensors-python"
174
- version = "0.5.2"
174
+ version = "0.5.3"
175
175
  dependencies = [
176
176
  "memmap2",
177
177
  "pyo3",
@@ -181,18 +181,18 @@ dependencies = [
181
181
 
182
182
  [[package]]
183
183
  name = "serde"
184
- version = "1.0.217"
184
+ version = "1.0.218"
185
185
  source = "registry+https://github.com/rust-lang/crates.io-index"
186
- checksum = "02fc4265df13d6fa1d00ecff087228cc0a2b5f3c0e87e258d8b94a156e984c70"
186
+ checksum = "e8dfc9d19bdbf6d17e22319da49161d5d0108e4188e8b680aef6299eed22df60"
187
187
  dependencies = [
188
188
  "serde_derive",
189
189
  ]
190
190
 
191
191
  [[package]]
192
192
  name = "serde_derive"
193
- version = "1.0.217"
193
+ version = "1.0.218"
194
194
  source = "registry+https://github.com/rust-lang/crates.io-index"
195
- checksum = "5a9bf7cf98d04a2b28aead066b7496853d4779c9cc183c440dbac457641e19a0"
195
+ checksum = "f09503e191f4e797cb8aac08e9a4a4695c5edf6a2e70e376d961ddd5c969f82b"
196
196
  dependencies = [
197
197
  "proc-macro2",
198
198
  "quote",
@@ -201,9 +201,9 @@ dependencies = [
201
201
 
202
202
  [[package]]
203
203
  name = "serde_json"
204
- version = "1.0.135"
204
+ version = "1.0.139"
205
205
  source = "registry+https://github.com/rust-lang/crates.io-index"
206
- checksum = "2b0d7ba2887406110130a978386c4e1befb98c674b4fba677954e4db976630d9"
206
+ checksum = "44f86c3acccc9c65b153fe1b85a3be07fe5515274ec9f0653b4a0875731c72a6"
207
207
  dependencies = [
208
208
  "itoa",
209
209
  "memchr",
@@ -213,9 +213,9 @@ dependencies = [
213
213
 
214
214
  [[package]]
215
215
  name = "syn"
216
- version = "2.0.95"
216
+ version = "2.0.98"
217
217
  source = "registry+https://github.com/rust-lang/crates.io-index"
218
- checksum = "46f71c0377baf4ef1cc3e3402ded576dccc315800fbc62dfc7fe04b009773b4a"
218
+ checksum = "36147f1a48ae0ec2b5b3bc5b537d267457555a10dc06f3dbc8cb11ba3006d3b1"
219
219
  dependencies = [
220
220
  "proc-macro2",
221
221
  "quote",
@@ -230,9 +230,9 @@ checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1"
230
230
 
231
231
  [[package]]
232
232
  name = "unicode-ident"
233
- version = "1.0.14"
233
+ version = "1.0.17"
234
234
  source = "registry+https://github.com/rust-lang/crates.io-index"
235
- checksum = "adb9e6ca4f869e1180728b7950e35922a7fc6397f7b641499e8f3ef06e50dc83"
235
+ checksum = "00e2473a93778eb0bad35909dff6a10d28e63f792f16ed15e404fca9d5eeedbe"
236
236
 
237
237
  [[package]]
238
238
  name = "unindent"
@@ -1,7 +1,8 @@
1
1
  [package]
2
2
  name = "safetensors-python"
3
- version = "0.5.2"
3
+ version = "0.5.3"
4
4
  edition = "2021"
5
+ rust-version = "1.74"
5
6
 
6
7
  # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
7
8
  [lib]
@@ -267,7 +267,8 @@ enum Device {
267
267
  Xpu(usize),
268
268
  Xla(usize),
269
269
  Mlu(usize),
270
- /// User didn't specify acceletor, torch
270
+ Hpu(usize),
271
+ /// User didn't specify accelerator, torch
271
272
  /// is responsible for choosing.
272
273
  Anonymous(usize),
273
274
  }
@@ -296,11 +297,13 @@ impl<'source> FromPyObject<'source> for Device {
296
297
  "xpu" => Ok(Device::Xpu(0)),
297
298
  "xla" => Ok(Device::Xla(0)),
298
299
  "mlu" => Ok(Device::Mlu(0)),
300
+ "hpu" => Ok(Device::Hpu(0)),
299
301
  name if name.starts_with("cuda:") => parse_device(name).map(Device::Cuda),
300
302
  name if name.starts_with("npu:") => parse_device(name).map(Device::Npu),
301
303
  name if name.starts_with("xpu:") => parse_device(name).map(Device::Xpu),
302
304
  name if name.starts_with("xla:") => parse_device(name).map(Device::Xla),
303
305
  name if name.starts_with("mlu:") => parse_device(name).map(Device::Mlu),
306
+ name if name.starts_with("hpu:") => parse_device(name).map(Device::Hpu),
304
307
  name => Err(SafetensorError::new_err(format!(
305
308
  "device {name} is invalid"
306
309
  ))),
@@ -327,6 +330,7 @@ impl<'py> IntoPyObject<'py> for Device {
327
330
  Device::Xpu(n) => format!("xpu:{n}").into_pyobject(py).map(|x| x.into_any()),
328
331
  Device::Xla(n) => format!("xla:{n}").into_pyobject(py).map(|x| x.into_any()),
329
332
  Device::Mlu(n) => format!("mlu:{n}").into_pyobject(py).map(|x| x.into_any()),
333
+ Device::Hpu(n) => format!("hpu:{n}").into_pyobject(py).map(|x| x.into_any()),
330
334
  Device::Anonymous(n) => n.into_pyobject(py).map(|x| x.into_any()),
331
335
  }
332
336
  }
@@ -795,8 +799,12 @@ struct Disp(Vec<TensorIndexer>);
795
799
  impl fmt::Display for Disp {
796
800
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
797
801
  write!(f, "[")?;
798
- for item in &self.0 {
799
- write!(f, "{item}")?;
802
+ for (i, item) in self.0.iter().enumerate() {
803
+ if i != self.0.len() - 1 {
804
+ write!(f, "{item}, ")?;
805
+ } else {
806
+ write!(f, "{item}")?;
807
+ }
800
808
  }
801
809
  write!(f, "]")
802
810
  }
@@ -170,6 +170,24 @@ class TorchTestCase(unittest.TestCase):
170
170
  for k, v in reloaded.items():
171
171
  self.assertTrue(torch.allclose(data[k], reloaded[k]))
172
172
 
173
+ def test_hpu(self):
174
+ # must be run to load torch with Intel Gaudi bindings
175
+ try:
176
+ import habana_frameworks.torch.core as htcore
177
+ except ImportError:
178
+ self.skipTest("HPU is not available")
179
+
180
+ data = {
181
+ "test1": torch.zeros((2, 2), dtype=torch.float32).to("hpu"),
182
+ "test2": torch.zeros((2, 2), dtype=torch.float16).to("hpu"),
183
+ }
184
+ local = "./tests/data/out_safe_pt_mmap_small_hpu.safetensors"
185
+ save_file(data, local)
186
+
187
+ reloaded = load_file(local, device="hpu")
188
+ for k, v in reloaded.items():
189
+ self.assertTrue(torch.allclose(data[k], reloaded[k]))
190
+
173
191
  @unittest.skipIf(not torch.cuda.is_available(), "Cuda is not available")
174
192
  def test_anonymous_accelerator(self):
175
193
  data = {
@@ -340,5 +340,19 @@ class ReadmeTestCase(unittest.TestCase):
340
340
  tensor = slice_[2:, 20]
341
341
  self.assertEqual(
342
342
  str(cm.exception),
343
- "Error during slicing [2:20] with shape [10, 5]: SliceOutOfRange { dim_index: 1, asked: 20, dim_size: 5 }",
343
+ "Error during slicing [2:, 20] with shape [10, 5]: SliceOutOfRange { dim_index: 1, asked: 20, dim_size: 5 }",
344
+ )
345
+
346
+ with self.assertRaises(SafetensorError) as cm:
347
+ tensor = slice_[:20]
348
+ self.assertEqual(
349
+ str(cm.exception),
350
+ "Error during slicing [:20] with shape [10, 5]: SliceOutOfRange { dim_index: 0, asked: 19, dim_size: 10 }",
351
+ )
352
+
353
+ with self.assertRaises(SafetensorError) as cm:
354
+ tensor = slice_[:, :20]
355
+ self.assertEqual(
356
+ str(cm.exception),
357
+ "Error during slicing [:, :20] with shape [10, 5]: SliceOutOfRange { dim_index: 1, asked: 19, dim_size: 5 }",
344
358
  )
@@ -1,7 +1,8 @@
1
1
  [package]
2
2
  name = "safetensors"
3
- version = "0.5.2"
3
+ version = "0.5.3"
4
4
  edition = "2021"
5
+ rust-version = "1.74"
5
6
  homepage = "https://github.com/huggingface/safetensors"
6
7
  repository = "https://github.com/huggingface/safetensors"
7
8
  documentation = "https://docs.rs/safetensors/"
@@ -7,6 +7,7 @@ use core::ops::{
7
7
 
8
8
  /// Error representing invalid slicing attempt
9
9
  #[derive(Debug)]
10
+ #[cfg_attr(test, derive(Eq, PartialEq))]
10
11
  pub enum InvalidSlice {
11
12
  /// When the client asked for more slices than the tensors has dimensions
12
13
  TooManySlices,
@@ -235,6 +236,7 @@ where
235
236
 
236
237
  /// Iterator used to return the bits of the overall tensor buffer
237
238
  /// when client asks for a slice of the original tensor.
239
+ #[cfg_attr(test, derive(Debug, Eq, PartialEq))]
238
240
  pub struct SliceIterator<'data> {
239
241
  view: &'data TensorView<'data>,
240
242
  indices: Vec<(usize, usize)>,
@@ -284,10 +286,15 @@ impl<'data> SliceIterator<'data> {
284
286
  }
285
287
  TensorIndexer::Select(s) => (*s, *s + 1),
286
288
  };
287
- if start >= shape && stop > shape {
289
+ if start >= shape || stop > shape {
290
+ let asked = if start >= shape {
291
+ start
292
+ } else {
293
+ stop.saturating_sub(1)
294
+ };
288
295
  return Err(InvalidSlice::SliceOutOfRange {
289
296
  dim_index: i,
290
- asked: stop.saturating_sub(1),
297
+ asked,
291
298
  dim_size: shape,
292
299
  });
293
300
  }
@@ -573,4 +580,54 @@ mod tests {
573
580
  assert_eq!(iterator.next(), Some(&data[12..16]));
574
581
  assert_eq!(iterator.next(), None);
575
582
  }
583
+
584
+ #[test]
585
+ fn test_invalid_range() {
586
+ let data: Vec<u8> = vec![0.0f32, 1.0, 2.0, 3.0, 4.0, 5.0]
587
+ .into_iter()
588
+ .flat_map(|f| f.to_le_bytes())
589
+ .collect();
590
+
591
+ let attn_0 = TensorView::new(Dtype::F32, vec![2, 3], &data).unwrap();
592
+
593
+ assert_eq!(
594
+ SliceIterator::new(
595
+ &attn_0,
596
+ &[
597
+ TensorIndexer::Select(1),
598
+ TensorIndexer::Narrow(Bound::Included(1), Bound::Excluded(4)),
599
+ ],
600
+ ),
601
+ Err(InvalidSlice::SliceOutOfRange {
602
+ asked: 3,
603
+ dim_index: 1,
604
+ dim_size: 3,
605
+ })
606
+ );
607
+ assert_eq!(
608
+ SliceIterator::new(
609
+ &attn_0,
610
+ &[
611
+ TensorIndexer::Select(1),
612
+ TensorIndexer::Narrow(Bound::Included(3), Bound::Excluded(2)),
613
+ ],
614
+ ),
615
+ Err(InvalidSlice::SliceOutOfRange {
616
+ asked: 3,
617
+ dim_index: 1,
618
+ dim_size: 3,
619
+ })
620
+ );
621
+ assert_eq!(
622
+ SliceIterator::new(
623
+ &attn_0,
624
+ &[
625
+ TensorIndexer::Select(1),
626
+ TensorIndexer::Select(1),
627
+ TensorIndexer::Select(1),
628
+ ],
629
+ ),
630
+ Err(InvalidSlice::TooManySlices)
631
+ );
632
+ }
576
633
  }
@@ -64,8 +64,12 @@ impl core::fmt::Display for SafeTensorError {
64
64
  }
65
65
  }
66
66
 
67
+ #[cfg(not(feature = "std"))]
67
68
  impl core::error::Error for SafeTensorError {}
68
69
 
70
+ #[cfg(feature = "std")]
71
+ impl std::error::Error for SafeTensorError {}
72
+
69
73
  struct PreparedData {
70
74
  n: u64,
71
75
  header_bytes: Vec<u8>,
File without changes