fn
compile_fused_training_step
→CompiledExecutable*CompiledExecutable * compile_fused_training_step(const TraceGraph & graph, const int & external_feeds, TensorId loss_id, const int & param_ids, const OptimizerSpec & opt_spec, const int)Compile a single MPSGraph executable that runs the full training step in one dispatch:
loss = loss_fn(model(x), target) grad_param[i] = ∂loss / ∂param[i] ← derived by MPSGraph new_param[i], new_state[i, k] = opt_update(param[i], grad_param[i], state[i, k], ...)
Output layout (in this order):
- loss ← single scalar/tensor produced by the forward.
- new_param[i] ← one per
param_ids[i]. - new_state[i, k] ← per (param, state-slot-index) — exactly the
state_buf_ids_per_paramlayout flattened row-major (i, k).
At run time the caller uses run_executable_inplace with the
output_targets = [loss_scratch, params (in order), state_bufs
(flattened)] so the new values overwrite the existing storage. The
loss is allocated fresh and returned to Python for logging.
Parameters
graphconst TraceGraph&Forward trace (model + loss). Same shape as
compile_trace_with_backward.external_feeds...id → impl map for every external feed in
graph. Must include every parameter id, every state buffer id, and every per-step scalar id.loss_id, param_ids :Forward-graph ids for the scalar loss and the parameters with respect to which the gradient is taken.
opt_spec :Hyperparameters + kind of the optimizer.
state_buf_ids_per_param :For each param i, the trace ids of its state buffers (in the order each optimizer expects: SGD = [momentum] or empty, Adam = [m, v]). All must already be in
external_feeds.scalar_input_ids :Per-step scalar feed ids in the order each optimizer expects (Adam / AdamW: [bias1, bias2]; SGD: empty). These are read at run time as 0-D float feeds.
error_msgstd::string*= NonePopulated on failure.
Returns
CompiledExecutable*Owns the compiled executable. output_ids = {loss_id}; grad_output_ids carries one id per (new_param, new_state) output in the order described above. Returns nullptr on abort.