Trait Optimizer 

pub trait Optimizer {
    // Required methods
    fn apply(
        &mut self,
        params: &[(&str, &Array)],
        grads: &[&Array],
    ) -> Result<Vec<Array>>;
    fn init_params(&mut self, params: &[(&str, &Array)]);
    fn visit_state(&self, v: &mut Visit<'_>);
    fn visit_state_mut(&mut self, v: &mut VisitMut<'_>);
    fn learning_rate(&self) -> Array;
    fn set_learning_rate(&mut self, lr: f32);
    fn clone_box(&self) -> Box<dyn Optimizer>;

    // Provided methods
    fn init<M: Tree>(&mut self, model: &M)
       where Self: Sized { ... }
    fn update<M: Tree>(&mut self, model: &mut M, grads: &M) -> Result<()>
       where Self: Sized { ... }
}

Updates a model’s trainable leaves from their gradients, as MLX’s Optimizer.

The trait is object-safe (Box<dyn Optimizer> is an optimizer and a tree too). Its state is reached only through Optimizer::visit_state and Optimizer::visit_state_mut (each optimizer has its own slot type).

Required Methods

fn apply( &mut self, params: &[(&str, &Array)], grads: &[&Array], ) -> Result<Vec<Array>>

One step over named parameters (the path keys each parameter’s state) and their gradients: the updated parameters, in order. Increments the step once.

fn init_params(&mut self, params: &[(&str, &Array)])

Creates zero state for params that have none (paths with state are kept).

fn visit_state(&self, v: &mut Visit<'_>)

Visits the state (step, learning_rate, slots) and the hyperparameters.

fn visit_state_mut(&mut self, v: &mut VisitMut<'_>)

Visits the state mutably.

fn learning_rate(&self) -> Array

The learning rate: the 0-D float32 array in the state.

fn set_learning_rate(&mut self, lr: f32)

Sets the learning rate array (in a compiled step it is an input: no retrace).

fn clone_box(&self) -> Box<dyn Optimizer>

A boxed copy (object-safe Clone).

Provided Methods

fn init<M: Tree>(&mut self, model: &M)

where Self: Sized,

Creates the state of every trainable leaf of model now (instead of at the first update), so a compiled step’s state structure is fixed from the start; also needed before loading a checkpoint into a fresh optimizer.

Examples found in repository

examples/mnist.rs (line 111)

fn main() -> Result<()> {
    let args: Vec<String> = std::env::args().skip(1).collect();
    let dir = args
        .iter()
        .find(|a| !a.starts_with("--"))
        .cloned()
        .unwrap_or_else(|| "data/mnist".into());
    let epochs: usize = args
        .iter()
        .position(|a| a == "--epochs")
        .and_then(|i| args.get(i + 1))
        .map_or(10, |v| v.parse().expect("--epochs takes a number"));
    let asynchronous = args.iter().any(|a| a == "--async");
    let compiled = args.iter().any(|a| a == "--compile");
    if args.iter().any(|a| a == "--cpu") {
        set_default_stream(Stream::default_for(Device::cpu()));
    }
    let (train_x, train_y, n) = load(&dir, "train").inspect_err(|_| {
        eprintln!("cannot read MNIST from {dir}: run `just mnist` (it downloads the data)");
    })?;
    let (test_x, test_y, n_test) = load(&dir, "t10k")?;
    let pixels = train_x.len() / n;
    let test_x = Array::from_slice(&test_x, [n_test, pixels]);
    let test_y = Array::from_slice(&test_y, [n_test]);

    let (batch_size, lr) = (256, 0.1);
    random::seed(0);
    let mut state = (Mlp::new(&[pixels as i32, 32, 32, 10]), Sgd::new(lr));
    state.1.init(&state.0); // the optimizer's state exists before the first step
    nuvec::eval(&state)?;
    // The training step, compiled with the model and the optimizer as its state.
    let step = nuvec::compile_mut(|(model, opt): &mut (Mlp, Sgd), (x, y): &(Array, Array)| {
        let loss_fn = |m: &Mlp| cross_entropy(&m.forward(x), y, Reduction::Mean);
        let (loss, grads) = value_and_grad(loss_fn)(model)?;
        opt.update(model, &grads)?;
        Ok(loss)
    });
    let mut shuffle = Shuffle(0);
    let device = nuvec::default_stream().device;
    let mut times = Vec::new();
    let mut accuracy = 0.0;
    for epoch in 0..epochs {
        let tic = Instant::now();
        let perm = shuffle.permutation(n);
        for ids in perm.chunks(batch_size) {
            let mut xb = Vec::with_capacity(ids.len() * pixels);
            for &i in ids {
                xb.extend_from_slice(&train_x[i * pixels..(i + 1) * pixels]);
            }
            let yb: Vec<u32> = ids.iter().map(|&i| train_y[i]).collect();
            let x = Array::from_slice(&xb, [ids.len(), pixels]);
            let y = Array::from_slice(&yb, [ids.len()]);
            let batch = (x, y);
            if compiled {
                let _loss = step(&mut state, &batch)?;
            } else {
                let (model, opt) = &mut state;
                let loss_fn =
                    |m: &Mlp| cross_entropy(&m.forward(&batch.0), &batch.1, Reduction::Mean);
                let (_, grads) = value_and_grad(loss_fn)(&*model)?;
                opt.update(model, &grads)?;
            }
            if asynchronous {
                nuvec::async_eval(&state.0)?;
            } else {
                nuvec::eval(&state.0)?;
            }
        }
        let predictions = state.0.forward(&test_x).argmax(1);
        let hits = predictions.equal(&test_y).astype(DType::F32);
        accuracy = hits.mean(..).item::<f32>()?;
        let seconds = tic.elapsed().as_secs_f64();
        times.push(seconds);
        println!("Epoch {epoch}: Test accuracy {accuracy:.3}, Time {seconds:.3} (s)");
    }
    let mut steady: Vec<f64> = times.iter().skip(1).copied().collect();
    steady.sort_by(f64::total_cmp);
    let median = steady.get(steady.len() / 2).copied().unwrap_or(times[0]);
    println!(
        "RESULT device={device} async={asynchronous} compiled={compiled} accuracy={accuracy:.4} epoch_seconds={median:.4}"
    );
    Ok(())
}

fn update<M: Tree>(&mut self, model: &mut M, grads: &M) -> Result<()>

where Self: Sized,

Applies one step to the trainable leaves of model from grads, a gradient tree of the model’s type with the model’s paths (else an error). Frozen parameters and buffers are left alone; their state is kept.

Examples found in repository

examples/mnist.rs (line 117)

fn main() -> Result<()> {
    let args: Vec<String> = std::env::args().skip(1).collect();
    let dir = args
        .iter()
        .find(|a| !a.starts_with("--"))
        .cloned()
        .unwrap_or_else(|| "data/mnist".into());
    let epochs: usize = args
        .iter()
        .position(|a| a == "--epochs")
        .and_then(|i| args.get(i + 1))
        .map_or(10, |v| v.parse().expect("--epochs takes a number"));
    let asynchronous = args.iter().any(|a| a == "--async");
    let compiled = args.iter().any(|a| a == "--compile");
    if args.iter().any(|a| a == "--cpu") {
        set_default_stream(Stream::default_for(Device::cpu()));
    }
    let (train_x, train_y, n) = load(&dir, "train").inspect_err(|_| {
        eprintln!("cannot read MNIST from {dir}: run `just mnist` (it downloads the data)");
    })?;
    let (test_x, test_y, n_test) = load(&dir, "t10k")?;
    let pixels = train_x.len() / n;
    let test_x = Array::from_slice(&test_x, [n_test, pixels]);
    let test_y = Array::from_slice(&test_y, [n_test]);

    let (batch_size, lr) = (256, 0.1);
    random::seed(0);
    let mut state = (Mlp::new(&[pixels as i32, 32, 32, 10]), Sgd::new(lr));
    state.1.init(&state.0); // the optimizer's state exists before the first step
    nuvec::eval(&state)?;
    // The training step, compiled with the model and the optimizer as its state.
    let step = nuvec::compile_mut(|(model, opt): &mut (Mlp, Sgd), (x, y): &(Array, Array)| {
        let loss_fn = |m: &Mlp| cross_entropy(&m.forward(x), y, Reduction::Mean);
        let (loss, grads) = value_and_grad(loss_fn)(model)?;
        opt.update(model, &grads)?;
        Ok(loss)
    });
    let mut shuffle = Shuffle(0);
    let device = nuvec::default_stream().device;
    let mut times = Vec::new();
    let mut accuracy = 0.0;
    for epoch in 0..epochs {
        let tic = Instant::now();
        let perm = shuffle.permutation(n);
        for ids in perm.chunks(batch_size) {
            let mut xb = Vec::with_capacity(ids.len() * pixels);
            for &i in ids {
                xb.extend_from_slice(&train_x[i * pixels..(i + 1) * pixels]);
            }
            let yb: Vec<u32> = ids.iter().map(|&i| train_y[i]).collect();
            let x = Array::from_slice(&xb, [ids.len(), pixels]);
            let y = Array::from_slice(&yb, [ids.len()]);
            let batch = (x, y);
            if compiled {
                let _loss = step(&mut state, &batch)?;
            } else {
                let (model, opt) = &mut state;
                let loss_fn =
                    |m: &Mlp| cross_entropy(&m.forward(&batch.0), &batch.1, Reduction::Mean);
                let (_, grads) = value_and_grad(loss_fn)(&*model)?;
                opt.update(model, &grads)?;
            }
            if asynchronous {
                nuvec::async_eval(&state.0)?;
            } else {
                nuvec::eval(&state.0)?;
            }
        }
        let predictions = state.0.forward(&test_x).argmax(1);
        let hits = predictions.equal(&test_y).astype(DType::F32);
        accuracy = hits.mean(..).item::<f32>()?;
        let seconds = tic.elapsed().as_secs_f64();
        times.push(seconds);
        println!("Epoch {epoch}: Test accuracy {accuracy:.3}, Time {seconds:.3} (s)");
    }
    let mut steady: Vec<f64> = times.iter().skip(1).copied().collect();
    steady.sort_by(f64::total_cmp);
    let median = steady.get(steady.len() / 2).copied().unwrap_or(times[0]);
    println!(
        "RESULT device={device} async={asynchronous} compiled={compiled} accuracy={accuracy:.4} epoch_seconds={median:.4}"
    );
    Ok(())
}

Trait Implementations

impl Clone for Box<dyn Optimizer>

fn clone(&self) -> Self

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for dyn Optimizer

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl Leaves for Box<dyn Optimizer>

A boxed optimizer is a tree through its state: the one hand-written tree impl of nuvec-nn (a trait object cannot be derived).

fn visit(&self, v: &mut Visit<'_>)

Reports every leaf (and child tree, enum variant, mode flag and static field) to v, in order.

impl Optimizer for Box<dyn Optimizer>

fn apply( &mut self, params: &[(&str, &Array)], grads: &[&Array], ) -> Result<Vec<Array>>

One step over named parameters (the path keys each parameter’s state) and their gradients: the updated parameters, in order. Increments the step once.

fn init_params(&mut self, params: &[(&str, &Array)])

Creates zero state for params that have none (paths with state are kept).

fn visit_state(&self, v: &mut Visit<'_>)

Visits the state (step, learning_rate, slots) and the hyperparameters.

fn visit_state_mut(&mut self, v: &mut VisitMut<'_>)

Visits the state mutably.

fn learning_rate(&self) -> Array

The learning rate: the 0-D float32 array in the state.

fn set_learning_rate(&mut self, lr: f32)

Sets the learning rate array (in a compiled step it is an input: no retrace).

fn clone_box(&self) -> Box<dyn Optimizer>

A boxed copy (object-safe Clone).

fn init<M: Tree>(&mut self, model: &M)

where Self: Sized,

Creates the state of every trainable leaf of model now (instead of at the first update), so a compiled step’s state structure is fixed from the start; also needed before loading a checkpoint into a fresh optimizer.

fn update<M: Tree>(&mut self, model: &mut M, grads: &M) -> Result<()>

where Self: Sized,

Applies one step to the trainable leaves of model from grads, a gradient tree of the model’s type with the model’s paths (else an error). Frozen parameters and buffers are left alone; their state is kept.

impl Tree for Box<dyn Optimizer>

fn visit_mut(&mut self, v: &mut VisitMut<'_>)

Reports every leaf (and child tree and mode flag) to v mutably, in the order of Leaves::visit.

fn for_each(&self, f: impl FnMut(&str, &Array))

Calls f(path, leaf) for every leaf, in order.

fn map(&self, f: impl FnMut(&str, &Array) -> Array) -> Self

The tree with every leaf replaced by f(path, leaf) (same structure; parameters keep their frozen flag).

fn leaves(&self) -> Vec<Array>

The leaves, in order.

fn paths(&self) -> Vec<String>

The leaves’ paths, in order.

fn flatten(&self) -> Vec<(String, Array)>

The leaves with their paths, in order (MLX’s tree_flatten).

fn unflatten( &self, leaves: impl IntoIterator<Item = Array>, ) -> Result<Self, Error>

A tree of this structure (self is the template) with leaves in order: exactly Tree::len of them, else an error naming both counts.

fn load_flat( &mut self, named: impl IntoIterator<Item = (String, Array)>, mode: LoadMode, ) -> Result<(), Error>

Replaces leaves by path. LoadMode::Strict requires exactly the tree’s paths with matching shapes; LoadMode::Partial assigns the matching names and ignores unknown ones (shapes must still match). Loaded dtypes are kept. On any problem the error lists every one (missing, unknown and mismatched names, sorted) and the tree is unchanged. Duplicate paths (in the tree or in named) are an error.

fn len(&self) -> usize

The number of leaves.

fn is_empty(&self) -> bool

Whether the tree has no leaves.

fn def(&self) -> TreeDef

The tree’s structure.

Dyn Compatibility

This trait is dyn compatible.

In older versions of Rust, dyn compatibility was called "object safety".

Implementations on Foreign Types

impl Optimizer for Box<dyn Optimizer>

fn apply( &mut self, params: &[(&str, &Array)], grads: &[&Array], ) -> Result<Vec<Array>>

fn init_params(&mut self, params: &[(&str, &Array)])

fn visit_state(&self, v: &mut Visit<'_>)

fn visit_state_mut(&mut self, v: &mut VisitMut<'_>)

fn learning_rate(&self) -> Array

fn set_learning_rate(&mut self, lr: f32)

fn clone_box(&self) -> Box<dyn Optimizer>

Implementors

impl Optimizer for Adam

impl Optimizer for AdamW

impl Optimizer for Lion

impl Optimizer for MultiOptimizer

impl Optimizer for Muon

impl Optimizer for Sgd