compile_train gives no way to get the shared variables its step writes

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

Start by tracing compile_train and the dict of shared-variable updates it assembles, then inspect how save_state consumes its step argument. Compare this with f.maker.expanded_inputs and the extra_outputs example. Done means callers can obtain the step's shared variables through the returned Function rather than rebuilding or reaching into internal inputs.