docs menu

Module Nx_io.Archive

Archives of entries.

An archive is an immutable collection of entries with distinct, non-empty names, whose arrays lie on one set. A structure (Nx.Ptree) enters an archive with Archive.of_value, which names each array by a prefix, a dot and the array's path (Nx.Ptree.Path.to_string), as "model.blocks.0.w". An empty prefix, the default, is left out with its dot, and so are the dot and the empty path of an array at the root.

Under a prefix p, a value owns p, every name that starts with p and a dot, and, for an empty p, every name. Archive.to_value reads the entries a value owns back into the shape and the placement of one the program already has, and ignores the others, so one file can store several values under distinct prefixes:

ocaml

let checkpoint = Nx_io.Archive.union [
  Nx_io.Archive.of_value ~prefix:"model" model params;
  Nx_io.Archive.of_value ~prefix:"opt" opt state ] in
Nx_io.save_safetensors path checkpoint;
let params =
  Nx_io.Archive.to_value ~prefix:"model" model ~like:params
    (Nx_io.load_safetensors path)

Code that reads a file another program wrote, an importer, asks for each entry by name, shape and dtype with Archive.tensor and Archive.float, and rearranges it with nx:

ocaml

let linear a ~inputs ~outputs name =
  (* The file stores the weight as [outputs; inputs]. *)
  Nx.matrix_transpose
    (Nx_io.Archive.float ~shape:[| outputs; inputs |] Nx.float32 name a)

Asking for an entry reads none of its bytes, except Archive.float when it converts. Only Archive.float converts an entry's dtype. A block-quantised or sub-byte entry is read as uint8 with Archive.tensor. An 8-bit float entry is read with Archive.tensor too, since its scales are other entries.

An archive is input data, so a mismatch between an archive and what is asked of it raises Failure, naming the entry.

#type 'd t

The type for archives of values on the set 'd.

#val of_list : (string * 'd Nx.packed) list -> 'd t

of_list entries is the archive of entries.

Raises Invalid_argument if a name is empty or given twice.

#val union : 'd t list -> 'd t

union ts is the archive of the entries of all ts.

Raises Invalid_argument if a name is in two of ts.

#val names : 'd t -> string list

names a is the names of a's entries, sorted.

#val find : string -> 'd t -> 'd Nx.packed option

find name a is a's entry name, if any.

val tensor : 
  shape:int array ->
  ('v, 's) Nx.dtype ->
  string ->
  'd t ->
  ('v, 's, 'd) Nx.t

tensor ~shape dt name a is a's entry name, as stored.

Raises Failure naming the entry if a has no entry name, or if its shape is not shape or its dtype is not dt, as in "Nx_io.Archive.tensor: h.0.w: shape [3] in the archive, [2; 3] asked for".

val float : 
  shape:int array ->
  (float, 's) Nx.dtype ->
  string ->
  'd t ->
  (float, 's, 'd) Nx.t

float ~shape dt name a is a's entry name at dt. It is the entry as stored when its dtype is dt, and otherwise its cast (Nx.cast), which allocates and reads every element. It converts between float16, bfloat16, float32 and float64. Each call that casts makes a new value, so an importer calls float once for two tied weights.

Raises Invalid_argument if dt is not one of the four. Raises Failure naming the entry if a has no entry name, if its shape is not shape, or if its dtype is not one of the four.

#val of_value : ?prefix:string -> ('s, 'd) Nx.Ptree.t -> 's -> 'd t

of_value ~prefix s x is the archive with one entry for each array s walks in x, fixed arrays included, named under prefix as the module's preamble states. The entries are x's arrays, and nothing is copied.

Raises Invalid_argument if two arrays of x have one name, as in "Nx_io.Archive.of_value: w: two arrays have this name", or if a name is empty, as that of an array at the root without a prefix.

#val to_value : ?prefix:string -> ('s, 'd) Nx.Ptree.t -> like:'s -> 'e t -> 's

to_value ~prefix s ~like a is like with each array replaced by a's entry of its name under prefix, placed where like's array lies (Nx.place). like gives the structure, with its list lengths, option presences and variant cases, and each array's dtype, shape and placement. Its arrays' elements are not read. For values x and y with equal Nx.Ptree.visits, dtypes, shapes and placements, to_value ~prefix s ~like:y (of_value ~prefix s x) is x.

Raises Failure, naming the entry, before any entry is placed:

if an entry like names is missing;
if its shape or dtype differs;
if an entry like owns is named by none of its arrays.

Raises Invalid_argument as of_value does for like's names.