fused_step(model: Module, loss_fn: Callable[..., Tensor], optimizer: Optimizer, grad_scaler: GradScaler | None = None)Return a callable that runs one fused training step.
The first call traces model(*x); loss_fn(out, *targets); optimizer._trace_update(...) once under a single Tracer,
plumbs the ghost-grad placeholders, and compiles the resulting
graph into one MPSGraphExecutable that runs forward +
loss + backward (via MPSGraph autodiff) + optimizer update in a
single submission. Subsequent calls reuse the cached executable.
Delegates the optimizer math to the matching
compile_optimizer subclass — so every Lucid eager optimizer
(all 13: SGD, Adam, AdamW, RMSprop, Adagrad, Adadelta, Adamax,
NAdam, SparseAdam, Rprop, ASGD, RAdam, LBFGS) compiles cleanly
into a fused step. LBFGS is the closure-less single-step variant
(per-element Barzilai-Borwein direction); full closure-driven
line search remains eager-only.
Parameters
modelModulenn.Module. Parameters from
optimizer.param_groups must reference these tensors.loss_fnCallableloss_fn(model_output, *targets). Captured by identity in
the trace, so passing a fresh closure each call defeats the
cache.optimizerOptimizerNotImplementedError from
compile_optimizer with the structural reason
(line-search / sign branch / per-step coefficient / …).lucid.amp.GradScaler contract entirely inside
the compiled executable:
- The loss is multiplied by the current
scaler._scalebefore the backward derivation, so MPSGraph autograd produces scaled gradients that won't underflow in F16. - Each scaled gradient is unscaled in F32 by
1/scaler._scalebefore the optimizer reads it. found_inf = OR(any(!isfinite(g_unscaled)))is computed across all params; each new_param / new_state is wrapped withwhere(found_inf, old, new)so an overflow step leaves params + state buffers unchanged.- After the executable runs,
found_infis read back to Python andscaler.update()is invoked — the scale halves on overflow, doubles aftergrowth_intervalclean steps.
step(*args) is always
the unscaled loss (the trace divides by scale on the
return path), matching the eager API exactly.Returns
Callable[..., Tensor]A callable step(*args) where args is
(model_input, *loss_targets). Returns the scalar loss
Tensor from the just-completed step (the parameter
and optimizer-state buffers have already been updated
in-place inside the executable).
Examples
>>> import lucid, lucid.nn as nn, lucid.nn.functional as F
>>> import lucid.optim as optim
>>> from lucid.compile import fused_step
>>> model = nn.Linear(8, 4).to('metal')
>>> opt = optim.Adam(model.parameters(), lr=1e-3)
>>> step = fused_step(model, F.cross_entropy, opt)
>>> for batch in batches:
... loss = step(batch.x, batch.target) # one executable submission
... # opt.step() is implicit — parameters already updated in-placeSee Also
- lucid.compile.make_step—forward+backward only, no optimizer fusion
(lets the caller drive the optimizer in Python).- lucid.compile.compile_optimizer—the underlying optimizer-side
translator that produces the update graph.