safetensors 0.5.1__tar.gz → 0.5.3__tar.gz

{safetensors-0.5.1 → safetensors-0.5.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: safetensors
-Version: 0.5.1
+Version: 0.5.3
 Classifier: Development Status :: 5 - Production/Stable
 Classifier: Intended Audience :: Developers
 Classifier: Intended Audience :: Education
@@ -14,31 +14,31 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Classifier: Typing :: Typed
-Requires-Dist: numpy >=1.21.6 ; extra == 'numpy'
+Requires-Dist: numpy>=1.21.6 ; extra == 'numpy'
 Requires-Dist: safetensors[numpy] ; extra == 'torch'
-Requires-Dist: torch >=1.10 ; extra == 'torch'
+Requires-Dist: torch>=1.10 ; extra == 'torch'
 Requires-Dist: safetensors[numpy] ; extra == 'tensorflow'
-Requires-Dist: tensorflow >=2.11.0 ; extra == 'tensorflow'
+Requires-Dist: tensorflow>=2.11.0 ; extra == 'tensorflow'
 Requires-Dist: safetensors[numpy] ; extra == 'pinned-tf'
-Requires-Dist: tensorflow ==2.18.0 ; extra == 'pinned-tf'
+Requires-Dist: tensorflow==2.18.0 ; extra == 'pinned-tf'
 Requires-Dist: safetensors[numpy] ; extra == 'jax'
-Requires-Dist: flax >=0.6.3 ; extra == 'jax'
-Requires-Dist: jax >=0.3.25 ; extra == 'jax'
-Requires-Dist: jaxlib >=0.3.25 ; extra == 'jax'
-Requires-Dist: mlx >=0.0.9 ; extra == 'mlx'
+Requires-Dist: flax>=0.6.3 ; extra == 'jax'
+Requires-Dist: jax>=0.3.25 ; extra == 'jax'
+Requires-Dist: jaxlib>=0.3.25 ; extra == 'jax'
+Requires-Dist: mlx>=0.0.9 ; extra == 'mlx'
 Requires-Dist: safetensors[numpy] ; extra == 'paddlepaddle'
-Requires-Dist: paddlepaddle >=2.4.1 ; extra == 'paddlepaddle'
-Requires-Dist: black ==22.3 ; extra == 'quality'
-Requires-Dist: click ==8.0.4 ; extra == 'quality'
-Requires-Dist: isort >=5.5.4 ; extra == 'quality'
-Requires-Dist: flake8 >=3.8.3 ; extra == 'quality'
+Requires-Dist: paddlepaddle>=2.4.1 ; extra == 'paddlepaddle'
+Requires-Dist: black==22.3 ; extra == 'quality'
+Requires-Dist: click==8.0.4 ; extra == 'quality'
+Requires-Dist: isort>=5.5.4 ; extra == 'quality'
+Requires-Dist: flake8>=3.8.3 ; extra == 'quality'
 Requires-Dist: safetensors[numpy] ; extra == 'testing'
-Requires-Dist: h5py >=3.7.0 ; extra == 'testing'
-Requires-Dist: huggingface-hub >=0.12.1 ; extra == 'testing'
-Requires-Dist: setuptools-rust >=1.5.2 ; extra == 'testing'
-Requires-Dist: pytest >=7.2.0 ; extra == 'testing'
-Requires-Dist: pytest-benchmark >=4.0.0 ; extra == 'testing'
-Requires-Dist: hypothesis >=6.70.2 ; extra == 'testing'
+Requires-Dist: h5py>=3.7.0 ; extra == 'testing'
+Requires-Dist: huggingface-hub>=0.12.1 ; extra == 'testing'
+Requires-Dist: setuptools-rust>=1.5.2 ; extra == 'testing'
+Requires-Dist: pytest>=7.2.0 ; extra == 'testing'
+Requires-Dist: pytest-benchmark>=4.0.0 ; extra == 'testing'
+Requires-Dist: hypothesis>=6.70.2 ; extra == 'testing'
 Requires-Dist: safetensors[torch] ; extra == 'all'
 Requires-Dist: safetensors[numpy] ; extra == 'all'
 Requires-Dist: safetensors[pinned-tf] ; extra == 'all'

{safetensors-0.5.1 → safetensors-0.5.3}/bindings/python/Cargo.lock

@@ -1,6 +1,6 @@
 # This file is automatically @generated by Cargo.
 # It is not intended for manual editing.
-version = 4
+version = 3
 
 [[package]]
 name = "autocfg"
@@ -34,9 +34,9 @@ checksum = "d75a2a4b1b190afb6f5425f10f6a8f959d2ea0b9c2b1d79553551850539e4674"
 
 [[package]]
 name = "libc"
-version = "0.2.169"
+version = "0.2.170"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b5aba8db14291edd000dfcc4d620c7ebfb122c613afb886ca8803fa4e128a20a"
+checksum = "875b3680cb2f8f71bdcf9a30f38d48282f5d3c95cbf9b3fa57269bb5d5c06828"
 
 [[package]]
 name = "memchr"
@@ -64,30 +64,30 @@ dependencies = [
 
 [[package]]
 name = "once_cell"
-version = "1.20.2"
+version = "1.20.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1261fe7e33c73b354eab43b1273a57c8f967d0391e80353e51f764ac02cf6775"
+checksum = "945462a4b81e43c4e3ba96bd7b49d834c6f61198356aa858733bc4acf3cbe62e"
 
 [[package]]
 name = "portable-atomic"
-version = "1.10.0"
+version = "1.11.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "280dc24453071f1b63954171985a0b0d30058d287960968b9b2aca264c8d4ee6"
+checksum = "350e9b48cbc6b0e028b0473b114454c6316e57336ee184ceab6e53f72c178b3e"
 
 [[package]]
 name = "proc-macro2"
-version = "1.0.92"
+version = "1.0.93"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "37d3544b3f2748c54e147655edb5025752e2303145b5aefb3c3ea2c78b973bb0"
+checksum = "60946a68e5f9d28b0dc1c21bb8a97ee7d018a8b322fa57838ba31cc878e22d99"
 dependencies = [
  "unicode-ident",
 ]
 
 [[package]]
 name = "pyo3"
-version = "0.23.3"
+version = "0.23.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e484fd2c8b4cb67ab05a318f1fd6fa8f199fcc30819f08f07d200809dba26c15"
+checksum = "7778bffd85cf38175ac1f545509665d0b9b92a198ca7941f131f85f7a4f9a872"
 dependencies = [
  "cfg-if",
  "indoc",
@@ -103,9 +103,9 @@ dependencies = [
 
 [[package]]
 name = "pyo3-build-config"
-version = "0.23.3"
+version = "0.23.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "dc0e0469a84f208e20044b98965e1561028180219e35352a2afaf2b942beff3b"
+checksum = "94f6cbe86ef3bf18998d9df6e0f3fc1050a8c5efa409bf712e661a4366e010fb"
 dependencies = [
  "once_cell",
  "target-lexicon",
@@ -113,9 +113,9 @@ dependencies = [
 
 [[package]]
 name = "pyo3-ffi"
-version = "0.23.3"
+version = "0.23.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "eb1547a7f9966f6f1a0f0227564a9945fe36b90da5a93b3933fc3dc03fae372d"
+checksum = "e9f1b4c431c0bb1c8fb0a338709859eed0d030ff6daa34368d3b152a63dfdd8d"
 dependencies = [
  "libc",
  "pyo3-build-config",
@@ -123,9 +123,9 @@ dependencies = [
 
 [[package]]
 name = "pyo3-macros"
-version = "0.23.3"
+version = "0.23.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fdb6da8ec6fa5cedd1626c886fc8749bdcbb09424a86461eb8cdf096b7c33257"
+checksum = "fbc2201328f63c4710f68abdf653c89d8dbc2858b88c5d88b0ff38a75288a9da"
 dependencies = [
  "proc-macro2",
  "pyo3-macros-backend",
@@ -135,9 +135,9 @@ dependencies = [
 
 [[package]]
 name = "pyo3-macros-backend"
-version = "0.23.3"
+version = "0.23.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "38a385202ff5a92791168b1136afae5059d3ac118457bb7bc304c197c2d33e7d"
+checksum = "fca6726ad0f3da9c9de093d6f116a93c1a38e417ed73bf138472cf4064f72028"
 dependencies = [
  "heck",
  "proc-macro2",
@@ -157,13 +157,13 @@ dependencies = [
 
 [[package]]
 name = "ryu"
-version = "1.0.18"
+version = "1.0.19"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f3cb5ba0dc43242ce17de99c180e96db90b235b8a9fdc9543c96d2209116bd9f"
+checksum = "6ea1a2d0a644769cc99faa24c3ad26b379b786fe7c36fd3c546254801650e6dd"
 
 [[package]]
 name = "safetensors"
-version = "0.5.1"
+version = "0.5.3"
 dependencies = [
  "serde",
  "serde_json",
@@ -171,7 +171,7 @@ dependencies = [
 
 [[package]]
 name = "safetensors-python"
-version = "0.5.1"
+version = "0.5.3"
 dependencies = [
  "memmap2",
  "pyo3",
@@ -181,18 +181,18 @@ dependencies = [
 
 [[package]]
 name = "serde"
-version = "1.0.217"
+version = "1.0.218"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "02fc4265df13d6fa1d00ecff087228cc0a2b5f3c0e87e258d8b94a156e984c70"
+checksum = "e8dfc9d19bdbf6d17e22319da49161d5d0108e4188e8b680aef6299eed22df60"
 dependencies = [
  "serde_derive",
 ]
 
 [[package]]
 name = "serde_derive"
-version = "1.0.217"
+version = "1.0.218"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5a9bf7cf98d04a2b28aead066b7496853d4779c9cc183c440dbac457641e19a0"
+checksum = "f09503e191f4e797cb8aac08e9a4a4695c5edf6a2e70e376d961ddd5c969f82b"
 dependencies = [
  "proc-macro2",
  "quote",
@@ -201,9 +201,9 @@ dependencies = [
 
 [[package]]
 name = "serde_json"
-version = "1.0.134"
+version = "1.0.139"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d00f4175c42ee48b15416f6193a959ba3a0d67fc699a0db9ad12df9f83991c7d"
+checksum = "44f86c3acccc9c65b153fe1b85a3be07fe5515274ec9f0653b4a0875731c72a6"
 dependencies = [
  "itoa",
  "memchr",
@@ -213,9 +213,9 @@ dependencies = [
 
 [[package]]
 name = "syn"
-version = "2.0.94"
+version = "2.0.98"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "987bc0be1cdea8b10216bd06e2ca407d40b9543468fafd3ddfb02f36e77f71f3"
+checksum = "36147f1a48ae0ec2b5b3bc5b537d267457555a10dc06f3dbc8cb11ba3006d3b1"
 dependencies = [
  "proc-macro2",
  "quote",
@@ -230,9 +230,9 @@ checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1"
 
 [[package]]
 name = "unicode-ident"
-version = "1.0.14"
+version = "1.0.17"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "adb9e6ca4f869e1180728b7950e35922a7fc6397f7b641499e8f3ef06e50dc83"
+checksum = "00e2473a93778eb0bad35909dff6a10d28e63f792f16ed15e404fca9d5eeedbe"
 
 [[package]]
 name = "unindent"

{safetensors-0.5.1 → safetensors-0.5.3}/bindings/python/Cargo.toml

@@ -1,7 +1,8 @@
 [package]
 name = "safetensors-python"
-version = "0.5.1"
+version = "0.5.3"
 edition = "2021"
+rust-version = "1.74"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 [lib]

{safetensors-0.5.1 → safetensors-0.5.3/bindings/python}/py_src/safetensors/__init__.pyi

@@ -69,7 +69,7 @@ class safe_open:
             The device on which you want the tensors.
     """
 
-    def __init__(filename, framework, device=...):
+    def __init__(self, filename, framework, device=...):
         pass
     def __enter__(self):
         """

{safetensors-0.5.1 → safetensors-0.5.3}/bindings/python/src/lib.rs

@@ -267,7 +267,8 @@ enum Device {
     Xpu(usize),
     Xla(usize),
     Mlu(usize),
-    /// User didn't specify acceletor, torch
+    Hpu(usize),
+    /// User didn't specify accelerator, torch
     /// is responsible for choosing.
     Anonymous(usize),
 }
@@ -296,11 +297,13 @@ impl<'source> FromPyObject<'source> for Device {
                 "xpu" => Ok(Device::Xpu(0)),
                 "xla" => Ok(Device::Xla(0)),
                 "mlu" => Ok(Device::Mlu(0)),
+                "hpu" => Ok(Device::Hpu(0)),
                 name if name.starts_with("cuda:") => parse_device(name).map(Device::Cuda),
                 name if name.starts_with("npu:") => parse_device(name).map(Device::Npu),
                 name if name.starts_with("xpu:") => parse_device(name).map(Device::Xpu),
                 name if name.starts_with("xla:") => parse_device(name).map(Device::Xla),
                 name if name.starts_with("mlu:") => parse_device(name).map(Device::Mlu),
+                name if name.starts_with("hpu:") => parse_device(name).map(Device::Hpu),
                 name => Err(SafetensorError::new_err(format!(
                     "device {name} is invalid"
                 ))),
@@ -327,6 +330,7 @@ impl<'py> IntoPyObject<'py> for Device {
             Device::Xpu(n) => format!("xpu:{n}").into_pyobject(py).map(|x| x.into_any()),
             Device::Xla(n) => format!("xla:{n}").into_pyobject(py).map(|x| x.into_any()),
             Device::Mlu(n) => format!("mlu:{n}").into_pyobject(py).map(|x| x.into_any()),
+            Device::Hpu(n) => format!("hpu:{n}").into_pyobject(py).map(|x| x.into_any()),
             Device::Anonymous(n) => n.into_pyobject(py).map(|x| x.into_any()),
         }
     }
@@ -795,8 +799,12 @@ struct Disp(Vec<TensorIndexer>);
 impl fmt::Display for Disp {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         write!(f, "[")?;
-        for item in &self.0 {
-            write!(f, "{item}")?;
+        for (i, item) in self.0.iter().enumerate() {
+            if i != self.0.len() - 1 {
+                write!(f, "{item}, ")?;
+            } else {
+                write!(f, "{item}")?;
+            }
         }
         write!(f, "]")
     }

{safetensors-0.5.1 → safetensors-0.5.3}/bindings/python/stub.py

@@ -42,10 +42,7 @@ def fn_predicate(obj):
         return (
             obj.__doc__
             and obj.__text_signature__
-            and (
-                not obj.__name__.startswith("_")
-                or obj.__name__ in {"__enter__", "__exit__"}
-            )
+            and (not obj.__name__.startswith("_") or obj.__name__ in {"__enter__", "__exit__"})
         )
     if inspect.isgetsetdescriptor(obj):
         return obj.__doc__ and not obj.__name__.startswith("_")
@@ -81,15 +78,14 @@ def pyi_file(obj, indent=""):
 
         body = ""
         if obj.__doc__:
-            body += (
-                f'{indent}"""\n{indent}{do_indent(obj.__doc__, indent)}\n{indent}"""\n'
-            )
+            body += f'{indent}"""\n{indent}{do_indent(obj.__doc__, indent)}\n{indent}"""\n'
 
         fns = inspect.getmembers(obj, fn_predicate)
 
         # Init
         if obj.__text_signature__:
-            body += f"{indent}def __init__{obj.__text_signature__}:\n"
+            signature = obj.__text_signature__.replace("(", "(self, ")
+            body += f"{indent}def __init__{signature}:\n"
             body += f"{indent+INDENT}pass\n"
             body += "\n"
 
@@ -146,11 +142,7 @@ def do_black(content, is_pyi):
 
 
 def write(module, directory, origin, check=False):
-    submodules = [
-        (name, member)
-        for name, member in inspect.getmembers(module)
-        if inspect.ismodule(member)
-    ]
+    submodules = [(name, member) for name, member in inspect.getmembers(module) if inspect.ismodule(member)]
 
     filename = os.path.join(directory, "__init__.pyi")
     pyi_content = pyi_file(module)
@@ -159,9 +151,7 @@ def write(module, directory, origin, check=False):
     if check:
         with open(filename, "r") as f:
             data = f.read()
-            assert (
-                data == pyi_content
-            ), f"The content of {filename} seems outdated, please run `python stub.py`"
+            assert data == pyi_content, f"The content of {filename} seems outdated, please run `python stub.py`"
     else:
         with open(filename, "w") as f:
             f.write(pyi_content)
@@ -184,9 +174,7 @@ def write(module, directory, origin, check=False):
         if check:
             with open(filename, "r") as f:
                 data = f.read()
-                assert (
-                    data == py_content
-                ), f"The content of {filename} seems outdated, please run `python stub.py`"
+                assert data == py_content, f"The content of {filename} seems outdated, please run `python stub.py`"
         else:
             with open(filename, "w") as f:
                 f.write(py_content)

{safetensors-0.5.1 → safetensors-0.5.3}/bindings/python/tests/test_pt_comparison.py

@@ -170,6 +170,24 @@ class TorchTestCase(unittest.TestCase):
         for k, v in reloaded.items():
             self.assertTrue(torch.allclose(data[k], reloaded[k]))
 
+    def test_hpu(self):
+        # must be run to load torch with Intel Gaudi bindings
+        try:
+            import habana_frameworks.torch.core as htcore
+        except ImportError:
+            self.skipTest("HPU is not available")
+
+        data = {
+            "test1": torch.zeros((2, 2), dtype=torch.float32).to("hpu"),
+            "test2": torch.zeros((2, 2), dtype=torch.float16).to("hpu"),
+        }
+        local = "./tests/data/out_safe_pt_mmap_small_hpu.safetensors"
+        save_file(data, local)
+
+        reloaded = load_file(local, device="hpu")
+        for k, v in reloaded.items():
+            self.assertTrue(torch.allclose(data[k], reloaded[k]))
+
     @unittest.skipIf(not torch.cuda.is_available(), "Cuda is not available")
     def test_anonymous_accelerator(self):
         data = {

{safetensors-0.5.1 → safetensors-0.5.3}/bindings/python/tests/test_simple.py

@@ -340,5 +340,19 @@ class ReadmeTestCase(unittest.TestCase):
                 tensor = slice_[2:, 20]
             self.assertEqual(
                 str(cm.exception),
-                "Error during slicing [2:20] with shape [10, 5]:  SliceOutOfRange { dim_index: 1, asked: 20, dim_size: 5 }",
+                "Error during slicing [2:, 20] with shape [10, 5]:  SliceOutOfRange { dim_index: 1, asked: 20, dim_size: 5 }",
+            )
+
+            with self.assertRaises(SafetensorError) as cm:
+                tensor = slice_[:20]
+            self.assertEqual(
+                str(cm.exception),
+                "Error during slicing [:20] with shape [10, 5]:  SliceOutOfRange { dim_index: 0, asked: 19, dim_size: 10 }",
+            )
+
+            with self.assertRaises(SafetensorError) as cm:
+                tensor = slice_[:, :20]
+            self.assertEqual(
+                str(cm.exception),
+                "Error during slicing [:, :20] with shape [10, 5]:  SliceOutOfRange { dim_index: 1, asked: 19, dim_size: 5 }",
             )

{safetensors-0.5.1/bindings/python → safetensors-0.5.3}/py_src/safetensors/__init__.pyi

@@ -69,7 +69,7 @@ class safe_open:
             The device on which you want the tensors.
     """
 
-    def __init__(filename, framework, device=...):
+    def __init__(self, filename, framework, device=...):
         pass
     def __enter__(self):
         """

{safetensors-0.5.1 → safetensors-0.5.3}/safetensors/Cargo.toml

@@ -1,7 +1,8 @@
 [package]
 name = "safetensors"
-version = "0.5.1"
+version = "0.5.3"
 edition = "2021"
+rust-version = "1.74"
 homepage = "https://github.com/huggingface/safetensors"
 repository = "https://github.com/huggingface/safetensors"
 documentation = "https://docs.rs/safetensors/"
@@ -21,14 +22,20 @@ exclude = [ "rust-toolchain", "target/*", "Cargo.lock"]
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-serde = {version = "1.0", features = ["derive"]}
-serde_json = "1.0"
+hashbrown = { version = "0.15.2", features = ["serde"], optional = true }
+serde = { version = "1.0", default-features = false, features = ["derive"] }
+serde_json = { version = "1.0", default-features = false }
 
 [dev-dependencies]
 criterion = "0.5"
 memmap2 = "0.9"
 proptest = "1.4"
 
+[features]
+default = ["std"]
+std = ["serde/default", "serde_json/default"]
+alloc = ["serde/alloc", "serde_json/alloc", "hashbrown"]
+
 [[bench]]
 name = "benchmark"
 harness = false

safetensors-0.5.3/safetensors/src/lib.rs

@@ -0,0 +1,43 @@
+#![deny(missing_docs)]
+#![doc = include_str!("../README.md")]
+#![cfg_attr(not(feature = "std"), no_std)]
+pub mod slice;
+pub mod tensor;
+/// serialize_to_file only valid in std
+#[cfg(feature = "std")]
+pub use tensor::serialize_to_file;
+pub use tensor::{serialize, Dtype, SafeTensorError, SafeTensors, View};
+
+#[cfg(feature = "alloc")]
+#[macro_use]
+extern crate alloc;
+
+#[cfg(all(feature = "std", feature = "alloc"))]
+compile_error!("must choose either the `std` or `alloc` feature, but not both.");
+#[cfg(all(not(feature = "std"), not(feature = "alloc")))]
+compile_error!("must choose either the `std` or `alloc` feature");
+
+/// A facade around all the types we need from the `std`, `core`, and `alloc`
+/// crates. This avoids elaborate import wrangling having to happen in every
+/// module.
+mod lib {
+    #[cfg(not(feature = "std"))]
+    mod no_stds {
+        pub use alloc::borrow::Cow;
+        pub use alloc::string::{String, ToString};
+        pub use alloc::vec::Vec;
+        pub use hashbrown::HashMap;
+    }
+    #[cfg(feature = "std")]
+    mod stds {
+        pub use std::borrow::Cow;
+        pub use std::collections::HashMap;
+        pub use std::string::{String, ToString};
+        pub use std::vec::Vec;
+    }
+    /// choose std or no_std to export by feature flag
+    #[cfg(not(feature = "std"))]
+    pub use no_stds::*;
+    #[cfg(feature = "std")]
+    pub use stds::*;
+}

{safetensors-0.5.1 → safetensors-0.5.3}/safetensors/src/slice.rs

@@ -1,12 +1,13 @@
 //! Module handling lazy loading via iterating on slices on the original buffer.
+use crate::lib::{String, ToString, Vec};
 use crate::tensor::TensorView;
-use std::fmt;
-use std::ops::{
+use core::ops::{
     Bound, Range, RangeBounds, RangeFrom, RangeFull, RangeInclusive, RangeTo, RangeToInclusive,
 };
 
 /// Error representing invalid slicing attempt
 #[derive(Debug)]
+#[cfg_attr(test, derive(Eq, PartialEq))]
 pub enum InvalidSlice {
     /// When the client asked for more slices than the tensors has dimensions
     TooManySlices,
@@ -40,8 +41,8 @@ fn display_bound(bound: &Bound<usize>) -> String {
 }
 
 /// Intended for Python users mostly or at least for its conventions
-impl fmt::Display for TensorIndexer {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+impl core::fmt::Display for TensorIndexer {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
         match self {
             TensorIndexer::Select(n) => {
                 write!(f, "{n}")
@@ -77,7 +78,7 @@ macro_rules! impl_from_range {
     ($range_type:ty) => {
         impl From<$range_type> for TensorIndexer {
             fn from(range: $range_type) -> Self {
-                use std::ops::Bound::*;
+                use core::ops::Bound::*;
 
                 let start = match range.start_bound() {
                     Included(idx) => Included(*idx),
@@ -235,6 +236,7 @@ where
 
 /// Iterator used to return the bits of the overall tensor buffer
 /// when client asks for a slice of the original tensor.
+#[cfg_attr(test, derive(Debug, Eq, PartialEq))]
 pub struct SliceIterator<'data> {
     view: &'data TensorView<'data>,
     indices: Vec<(usize, usize)>,
@@ -284,10 +286,15 @@ impl<'data> SliceIterator<'data> {
                     }
                     TensorIndexer::Select(s) => (*s, *s + 1),
                 };
-                if start >= shape && stop > shape {
+                if start >= shape || stop > shape {
+                    let asked = if start >= shape {
+                        start
+                    } else {
+                        stop.saturating_sub(1)
+                    };
                     return Err(InvalidSlice::SliceOutOfRange {
                         dim_index: i,
-                        asked: stop.saturating_sub(1),
+                        asked,
                         dim_size: shape,
                     });
                 }
@@ -573,4 +580,54 @@ mod tests {
         assert_eq!(iterator.next(), Some(&data[12..16]));
         assert_eq!(iterator.next(), None);
     }
+
+    #[test]
+    fn test_invalid_range() {
+        let data: Vec<u8> = vec![0.0f32, 1.0, 2.0, 3.0, 4.0, 5.0]
+            .into_iter()
+            .flat_map(|f| f.to_le_bytes())
+            .collect();
+
+        let attn_0 = TensorView::new(Dtype::F32, vec![2, 3], &data).unwrap();
+
+        assert_eq!(
+            SliceIterator::new(
+                &attn_0,
+                &[
+                    TensorIndexer::Select(1),
+                    TensorIndexer::Narrow(Bound::Included(1), Bound::Excluded(4)),
+                ],
+            ),
+            Err(InvalidSlice::SliceOutOfRange {
+                asked: 3,
+                dim_index: 1,
+                dim_size: 3,
+            })
+        );
+        assert_eq!(
+            SliceIterator::new(
+                &attn_0,
+                &[
+                    TensorIndexer::Select(1),
+                    TensorIndexer::Narrow(Bound::Included(3), Bound::Excluded(2)),
+                ],
+            ),
+            Err(InvalidSlice::SliceOutOfRange {
+                asked: 3,
+                dim_index: 1,
+                dim_size: 3,
+            })
+        );
+        assert_eq!(
+            SliceIterator::new(
+                &attn_0,
+                &[
+                    TensorIndexer::Select(1),
+                    TensorIndexer::Select(1),
+                    TensorIndexer::Select(1),
+                ],
+            ),
+            Err(InvalidSlice::TooManySlices)
+        );
+    }
 }

{safetensors-0.5.1 → safetensors-0.5.3}/safetensors/src/tensor.rs

@@ -1,11 +1,9 @@
 //! Module Containing the most important structures
+use crate::lib::{Cow, HashMap, String, ToString, Vec};
 use crate::slice::{InvalidSlice, SliceIterator, TensorIndexer};
 use serde::{ser::SerializeMap, Deserialize, Deserializer, Serialize, Serializer};
-use std::borrow::Cow;
-use std::collections::HashMap;
-use std::fs::File;
-use std::io::{BufWriter, Write};
-use std::path::Path;
+#[cfg(feature = "std")]
+use std::io::Write;
 
 const MAX_HEADER_SIZE: usize = 100_000_000;
 
@@ -32,6 +30,7 @@ pub enum SafeTensorError {
     /// The offsets declared for tensor with name `String` in the header are invalid
     InvalidOffset(String),
     /// IoError
+    #[cfg(feature = "std")]
     IoError(std::io::Error),
     /// JSON error
     JsonError(serde_json::Error),
@@ -46,6 +45,7 @@ pub enum SafeTensorError {
     ValidationOverflow,
 }
 
+#[cfg(feature = "std")]
 impl From<std::io::Error> for SafeTensorError {
     fn from(error: std::io::Error) -> SafeTensorError {
         SafeTensorError::IoError(error)
@@ -58,12 +58,16 @@ impl From<serde_json::Error> for SafeTensorError {
     }
 }
 
-impl std::fmt::Display for SafeTensorError {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+impl core::fmt::Display for SafeTensorError {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
         write!(f, "{self:?}")
     }
 }
 
+#[cfg(not(feature = "std"))]
+impl core::error::Error for SafeTensorError {}
+
+#[cfg(feature = "std")]
 impl std::error::Error for SafeTensorError {}
 
 struct PreparedData {
@@ -164,7 +168,7 @@ pub trait View {
     fn data_len(&self) -> usize;
 }
 
-fn prepare<S: AsRef<str> + Ord + std::fmt::Display, V: View, I: IntoIterator<Item = (S, V)>>(
+fn prepare<S: AsRef<str> + Ord + core::fmt::Display, V: View, I: IntoIterator<Item = (S, V)>>(
     data: I,
     data_info: &Option<HashMap<String, String>>,
     // ) -> Result<(Metadata, Vec<&'hash TensorView<'data>>, usize), SafeTensorError> {
@@ -212,7 +216,7 @@ fn prepare<S: AsRef<str> + Ord + std::fmt::Display, V: View, I: IntoIterator<Ite
 
 /// Serialize to an owned byte buffer the dictionnary of tensors.
 pub fn serialize<
-    S: AsRef<str> + Ord + std::fmt::Display,
+    S: AsRef<str> + Ord + core::fmt::Display,
     V: View,
     I: IntoIterator<Item = (S, V)>,
 >(
@@ -240,14 +244,15 @@ pub fn serialize<
 /// Serialize to a regular file the dictionnary of tensors.
 /// Writing directly to file reduces the need to allocate the whole amount to
 /// memory.
+#[cfg(feature = "std")]
 pub fn serialize_to_file<
-    S: AsRef<str> + Ord + std::fmt::Display,
+    S: AsRef<str> + Ord + core::fmt::Display,
     V: View,
     I: IntoIterator<Item = (S, V)>,
 >(
     data: I,
     data_info: &Option<HashMap<String, String>>,
-    filename: &Path,
+    filename: &std::path::Path,
 ) -> Result<(), SafeTensorError> {
     let (
         PreparedData {
@@ -255,7 +260,7 @@ pub fn serialize_to_file<
         },
         tensors,
     ) = prepare(data, data_info)?;
-    let mut f = BufWriter::new(File::create(filename)?);
+    let mut f = std::io::BufWriter::new(std::fs::File::create(filename)?);
     f.write_all(n.to_le_bytes().as_ref())?;
     f.write_all(&header_bytes)?;
     for tensor in tensors {
@@ -303,7 +308,7 @@ impl<'data> SafeTensors<'data> {
             return Err(SafeTensorError::InvalidHeaderLength);
         }
         let string =
-            std::str::from_utf8(&buffer[8..stop]).map_err(|_| SafeTensorError::InvalidHeader)?;
+            core::str::from_utf8(&buffer[8..stop]).map_err(|_| SafeTensorError::InvalidHeader)?;
         // Assert the string starts with {
         // NOTE: Add when we move to 0.4.0
         // if !string.starts_with('{') {
@@ -719,6 +724,9 @@ mod tests {
     use super::*;
     use crate::slice::IndexOp;
     use proptest::prelude::*;
+    #[cfg(not(feature = "std"))]
+    extern crate std;
+    use std::io::Write;
 
     const MAX_DIMENSION: usize = 8;
     const MAX_SIZE: usize = 8;
@@ -1021,10 +1029,13 @@ mod tests {
         std::fs::remove_file(&filename).unwrap();
 
         // File api
-        serialize_to_file(&metadata, &None, Path::new(&filename)).unwrap();
-        let raw = std::fs::read(&filename).unwrap();
-        let _deserialized = SafeTensors::deserialize(&raw).unwrap();
-        std::fs::remove_file(&filename).unwrap();
+        #[cfg(feature = "std")]
+        {
+            serialize_to_file(&metadata, &None, std::path::Path::new(&filename)).unwrap();
+            let raw = std::fs::read(&filename).unwrap();
+            let _deserialized = SafeTensors::deserialize(&raw).unwrap();
+            std::fs::remove_file(&filename).unwrap();
+        }
     }
 
     #[test]
@@ -1097,7 +1108,7 @@ mod tests {
         let n = serialized.len();
 
         let filename = "out.safetensors";
-        let mut f = BufWriter::new(File::create(filename).unwrap());
+        let mut f = std::io::BufWriter::new(std::fs::File::create(filename).unwrap());
         f.write_all(n.to_le_bytes().as_ref()).unwrap();
         f.write_all(serialized).unwrap();
         f.write_all(b"\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0").unwrap();

safetensors-0.5.1/safetensors/src/lib.rs

@@ -1,5 +0,0 @@
-#![deny(missing_docs)]
-#![doc = include_str!("../README.md")]
-pub mod slice;
-pub mod tensor;
-pub use tensor::{serialize, serialize_to_file, Dtype, SafeTensorError, SafeTensors, View};