package nx

  1. Overview
  2. Docs
package nx
Legend:
Page
Library
Module
Module type
Parameter
Class
Class type
Source

Module SafetensorsSource

Safetensors - Simple, safe way to store and distribute tensors

This is an OCaml port of the safetensors format, providing efficient serialization and deserialization of tensor data with metadata support.

Error Handling

Sourcetype safetensor_error =
  1. | Invalid_header of string
    (*

    Invalid UTF-8 in header

    *)
  2. | Invalid_header_start
    (*

    Invalid start character in header

    *)
  3. | Invalid_header_deserialization of string
    (*

    JSON parse error

    *)
  4. | Header_too_large
    (*

    Header exceeds maximum size

    *)
  5. | Header_too_small
    (*

    Header is too small

    *)
  6. | Invalid_header_length
    (*

    Invalid header length

    *)
  7. | Tensor_not_found of string
    (*

    Tensor with given name not found

    *)
  8. | Tensor_invalid_info
    (*

    Invalid shape, dtype, or offset for tensor

    *)
  9. | Invalid_offset of string
    (*

    Invalid offset for tensor

    *)
  10. | Io_error of string
    (*

    I/O error during file operations

    *)
  11. | Json_error of string
    (*

    JSON processing error

    *)
  12. | Invalid_tensor_view of string * int list * int
    (*

    Invalid tensor view creation

    *)
  13. | Metadata_incomplete_buffer
    (*

    Incomplete metadata, file not fully covered

    *)
  14. | Validation_overflow
    (*

    Overflow computing buffer size

    *)
  15. | Misaligned_slice
    (*

    Slice not aligned to byte boundary for sub-byte dtypes

    *)
Sourceval string_of_error : safetensor_error -> string

Convert an error to a human-readable string

Data Types

Sourcetype dtype =
  1. | BOOL
  2. | F4
    (*

    4-bit float

    *)
  3. | F6_E2M3
    (*

    6-bit float with 2 exponent bits, 3 mantissa bits

    *)
  4. | F6_E3M2
    (*

    6-bit float with 3 exponent bits, 2 mantissa bits

    *)
  5. | U8
    (*

    Unsigned 8-bit integer

    *)
  6. | I8
    (*

    Signed 8-bit integer

    *)
  7. | F8_E5M2
    (*

    8-bit float with 5 exponent bits, 2 mantissa bits

    *)
  8. | F8_E4M3
    (*

    8-bit float with 4 exponent bits, 3 mantissa bits

    *)
  9. | F8_E8M0
    (*

    8-bit float with 8 exponent bits, no mantissa

    *)
  10. | I16
    (*

    Signed 16-bit integer

    *)
  11. | U16
    (*

    Unsigned 16-bit integer

    *)
  12. | F16
    (*

    Half-precision float

    *)
  13. | BF16
    (*

    Brain float 16

    *)
  14. | I32
    (*

    Signed 32-bit integer

    *)
  15. | U32
    (*

    Unsigned 32-bit integer

    *)
  16. | F32
    (*

    Single-precision float

    *)
  17. | F64
    (*

    Double-precision float

    *)
  18. | I64
    (*

    Signed 64-bit integer

    *)
  19. | U64
    (*

    Unsigned 64-bit integer

    *)
Sourceval dtype_to_string : dtype -> string

Convert a dtype to its string representation

Sourceval dtype_of_string : string -> dtype option

Parse a dtype from its string representation

Sourceval bitsize : dtype -> int

Get the size in bits of one element of this dtype

Tensor Views

Sourcetype tensor_view = {
  1. dtype : dtype;
  2. shape : int list;
  3. data : string;
    (*

    Backing buffer

    *)
  4. offset : int;
    (*

    Byte offset into data

    *)
  5. length : int;
    (*

    Number of bytes

    *)
}

A view into tensor data without copying

Sourceval tensor_view_new : dtype:dtype -> shape:int list -> data:string -> (tensor_view, safetensor_error) result

Create a new tensor view. The data must exactly match the expected size based on dtype and shape.

SafeTensors Container

Sourcetype tensor_info = {
  1. dtype : dtype;
  2. shape : int list;
  3. data_offsets : int * int;
}
Sourcetype metadata = {
  1. metadata_kv : (string * string) list option;
  2. tensors : tensor_info array;
  3. index_map : (string, int) Hashtbl.t;
}

Metadata structure containing tensor information and optional key-value pairs

Sourcetype safetensors = {
  1. metadata : metadata;
  2. data : string;
}

The main container holding tensor metadata and data

Sourceval deserialize : string -> (safetensors, safetensor_error) result

Deserialize a safetensors buffer into a container. The buffer should contain the complete safetensors file content.

Sourceval serialize : (string * tensor_view) list -> (string * string) list option -> (string, safetensor_error) result

Serialize a list of named tensors with optional metadata key-value pairs into a safetensors buffer.

Sourceval serialize_to_file : (string * tensor_view) list -> (string * string) list option -> string -> (unit, safetensor_error) result

Serialize tensors directly to a file. This is more memory-efficient than serialize for large tensors.

Tensor Access

Get a specific tensor by name from the container

Sourceval tensors : safetensors -> (string * tensor_view) list

Get all tensors from the container as a list of name-view pairs

Sourceval iter : safetensors -> (string * tensor_view) list

Iterate over tensors in offset order (the order they appear in the file)

Sourceval names : safetensors -> string list

Get the names of all tensors in the container

Sourceval len : safetensors -> int

Get the number of tensors in the container

Sourceval is_empty : safetensors -> bool

Check if the container has no tensors

Slicing

Sourcetype bound =
  1. | Unbounded
  2. | Excluded of int
  3. | Included of int
Sourcetype tensor_indexer =
  1. | Select of int
  2. | Narrow of bound * bound
Sourcetype invalid_slice =
  1. | Too_many_slices
  2. | Slice_out_of_range of {
    1. dim_index : int;
    2. asked : int;
    3. dim_size : int;
    }
  3. | Misaligned_slices
Sourcetype slice_iterator = {
  1. view : tensor_view;
  2. mutable indices : (int * int) list;
  3. newshape : int list;
}
Sourcemodule Slice : sig ... end

Slicing support for lazy iteration over tensor data

Low-level Functions

Sourceval read_u64_le : string -> int -> int64

Read a little-endian 64-bit integer from a string at the given offset

Sourceval write_u64_le : Bytes.t -> int -> int64 -> unit

Write a little-endian 64-bit integer to a byte buffer at the given offset