save_state wants the list of shared variables one step writes and compile_train assembles exactly that dict and throws it away, so anything that checkpoints -- or that has to force MLX's lazy arrays each step -- hand-assembles the updates instead, or reaches into f.maker.expanded_inputs.
step = compile_train(loss, adam(1e-3), extra_outputs=[abs_residual])
save_state(step, "latest.safetensors") # want: a Function is enough to say what to save
save_statewants the list of shared variables one step writes andcompile_trainassembles exactly that dict and throws it away, so anything that checkpoints -- or that has to force MLX's lazy arrays each step -- hand-assembles the updates instead, or reaches intof.maker.expanded_inputs.