Function compile_mut 

pub fn compile_mut<'f, S, I, O, F, R>(
    f: F,
) -> impl Fn(&mut S, &I) -> Result<O> + 'f
where
    S: Tree,
    I: Tree,
    O: Tree,
    F: Fn(&mut S, &I) -> R + 'f,
    R: IntoResult<Ok = O>,

Compiles a stateful step, as MLX’s mx.compile(step, inputs=state, outputs=state): f(&mut state, &inputs) where state is a tree (typically (model, optimizer)).

use nuvec_compile::compile_mut;
use nuvec_core::{Array, Result, array, value_and_grad};

// Gradient descent on w for loss(w) = sum((w x - y)^2), compiled.
let step = compile_mut(|w: &mut Array, (x, y): &(Array, Array)| {
    let loss = |w: &Array| {
        let r = w.multiply(x).subtract(y);
        r.multiply(&r).sum(..)
    };
    let (l, g) = value_and_grad(loss)(w)?;
    *w = w.subtract(g.multiply(0.1));
    Ok(l)
});
let mut w = array![0.0];
for _ in 0..50 {
    step(&mut w, &(array![1.0], array![2.0]))?;
}
assert!((w.item::<f32>()? - 2.0).abs() < 1e-3);

• The state’s leaves (and the thread’s random key) are traced inputs; after the call the state is rebuilt from the tape’s outputs. So the step traces once, an optimizer’s learning-rate array set between calls is just a new input, and random draws inside (dropout) advance the thread’s key sequence exactly as an uncompiled call would.
• The tape is keyed by the state’s and the inputs’ structure: switching train/eval, freezing, changing a hashed hyperparameter retraces once.
• A step that changes the state’s structure (an optimizer creating its slots at the first update) retraces once; a structure that changes on two consecutive calls is an error (call opt.init(&model) first). Host fields mutated in f are hashed statics, so mutating one counts as a structure change; #[tree(skip, static_unchecked)] fields are baked into the trace.
• On error the state is unchanged. Called inside another transform, it runs f directly.