nuvec_io/
tree.rs

1//! Any tree of arrays as a safetensors file, by path: [`save`] and [`load_into`].
2
3use std::path::Path;
4
5use nuvec_core::tree::{LeafKind, Leaves, LoadMode, Tree, Visit, Visitor};
6use nuvec_core::{Array, DType, Error, Result};
7
8use crate::{load_safetensors, save_safetensors};
9
10/// Writes the leaves of `tree` as a safetensors file, each named by its path (MLX's
11/// `tree_flatten` names), evaluating them.
12///
13/// The file's metadata records `format = "nuvec-tree"`, the root type name and, when the
14/// tree holds 0-D leaves named `step` and `learning_rate` (an optimizer's state), their
15/// values (informational: the tensors are authoritative). Checkpointing a training run is
16/// `nuvec::save(path, &(model, opt))?`, resuming it
17/// `nuvec::load_into(path, &mut (model, opt))?` (after `opt.init(&model)`).
18///
19/// Duplicate paths are an error.
20pub fn save<T: Leaves + ?Sized>(path: impl AsRef<Path>, tree: &T) -> Result<()> {
21    struct Named(Vec<(String, Array)>);
22    impl Visitor for Named {
23        fn leaf(&mut self, path: &str, a: &Array, _: LeafKind, _: bool) {
24            self.0.push((path.to_owned(), a.clone()));
25        }
26    }
27    let mut named = Named(Vec::new());
28    tree.visit(&mut Visit::new(&mut named));
29    let named = named.0;
30    let mut seen = std::collections::HashSet::new();
31    for (p, _) in &named {
32        if !seen.insert(p.as_str()) {
33            return Err(Error::InvalidArgument(format!(
34                "save: the tree has the path {p:?} twice"
35            )));
36        }
37    }
38    let mut metadata = vec![
39        ("format".to_owned(), "nuvec-tree".to_owned()),
40        ("root".to_owned(), std::any::type_name::<T>().to_owned()),
41    ];
42    for key in ["step", "learning_rate"] {
43        let found = named
44            .iter()
45            .find(|(p, a)| a.ndim() == 0 && (p == key || p.ends_with(&format!(".{key}"))));
46        if let Some((_, a)) = found {
47            let value = match a.dtype() {
48                DType::F16 | DType::BF16 | DType::F32 | DType::F64 => {
49                    a.astype(DType::F64).item::<f64>()?.to_string()
50                }
51                _ => a.astype(DType::I64).item::<i64>()?.to_string(),
52            };
53            metadata.push((key.to_owned(), value));
54        }
55    }
56    let tensors: Vec<(&str, &Array)> = named.iter().map(|(p, a)| (p.as_str(), a)).collect();
57    let metadata: Vec<(&str, &str)> = metadata
58        .iter()
59        .map(|(k, v)| (k.as_str(), v.as_str()))
60        .collect();
61    save_safetensors(path, &tensors, &metadata)
62}
63
64/// Reads a file written by [`save`] into `tree` by path, strictly: every leaf must be in
65/// the file with its shape, and nothing else (loaded dtypes are kept). On any problem the
66/// error lists every one and `tree` is unchanged.
67pub fn load_into<T: Tree>(path: impl AsRef<Path>, tree: &mut T) -> Result<()> {
68    let file = load_safetensors(path)?;
69    tree.load_flat(file.tensors, LoadMode::Strict)
70        .map_err(|e| match e {
71            Error::InvalidArgument(m) => {
72                Error::InvalidArgument(m.replacen("load_flat", "load_into", 1))
73            }
74            e => e,
75        })
76}
nuvec_io/tree.rs

nuvec_io/
tree.rs
