Crate ndarray_safetensors

Source
Expand description

Serialize / deserialize Rust ndarray with Safetensors.

§Main APIs

§Demo

// Serailize ndarrays
let arr = array![[1.0, -1.0], [2.0, -2.0]];
let data = vec![("arr", TensorViewWithDataBuffer::new(&arr))];
let serialized_data = safetensors::serialize(data, &None).unwrap();
 
// Deserialize ndarrays
let tensors = safetensors::SafeTensors::deserialize(&serialized_data).unwrap();
let arrays = parse_tensors::<f64>(&tensors).unwrap();
let (name, array) = &arrays[0];
assert_eq!(name, "arr");
assert_eq!(array[[1,1]], -2.0);
 
// Deserialize with a wrong type hint
let parse_with_wrong_type = parse_tensors::<f32>(&tensors);
assert!(parse_with_wrong_type.is_err())

§Features

The following crate feature flags are available. They are configured in your Cargo.toml.

  • unsafe_copy: Allowing directly copying the bytes instead of parsing floats in serialization/deserialization for little-endian machines.
  • x86_sse: Using SSE instructions to optimize some conversions, it could be useful for fp16.

§License

Copyright (c) 2024, Mengxiao Lin. The crate is published under MIT License.

Structs§

TensorViewWithDataBuffer
A data structure like TensorView in safetensors, but it owns the data

Enums§

DeserializationError
Error to emit if the type doesn’t match for parsing a tensor view

Traits§

BFloat16ConversionSupportedElement
Element type traits that can be used to load/save Brain float point 16 numbers.
CommonSupportedElement
Element type traits for data types supported by both ndarray and safetensors
Float16ConversionSupportedElement
Element type traits that can be used to load/save IEEE 754 binary 16 floating point numbers

Functions§

parse_bf16_tensor_view_data
Deserialized a BF16 Safetensors View as a ndarray
parse_fp16_tensor_view_data
Deserialized a IEEE 754 FP16 Safetensors View as a ndarray.
parse_tensor_view_data
Deserialized a Safetensors View as a ndarray
parse_tensor_view_data_with_dimension
Deserialized a Safetensors View as a ndarray with a known dimension.
parse_tensors
Deserialize safetensors to ndarrays of the same element type and dimension type.