compiled_step(model: Module, x: Tensor, loss_fn: Callable[[Tensor], Tensor], use_grad_mode: bool = False)Run one compiled training step: forward + backward + grad assignment.
Parameters
modelnn.ModuleModule whose
forward(x) produces the model output.Input tensor (single forward pass).
loss_fncallableMaps
model(x) → scalar loss tensor. Typical: lambda y: F.cross_entropy(y, target).use_grad_modebool= FalseWhether to keep autograd's GradMode on during tracing. Default
False — the compiled step bypasses Lucid's autograd graph
entirely; we only need the trace hook to fire (it fires under
no_grad too because the kernel hooks run before the GradMode
short-circuit). Set True if you want the eager autograd
graph to also build (useful for verification).Returns
TensorThe loss tensor (compiled result). After this call,
param.grad is populated for every parameter of model.
The caller can then run an optimizer step.
Notes
The trace is captured fresh on every call; cache integration
(compile-or-cached) lives in Phase 1.4's CompiledModule.
Phase 1.3 acceptance: ResNet-18 + CE + SGD, 5 steps, param
drift < 5e-3 vs eager.
Examples
>>> import lucid, lucid.nn as nn, lucid.nn.functional as F
>>> import lucid.optim as optim
>>> from lucid.compile._entry.function import compiled_step
>>> model = nn.Linear(8, 4).to('metal')
>>> opt = optim.SGD(model.parameters(), lr=0.1)
>>> for batch in batches:
... opt.zero_grad()
... loss = compiled_step(model, batch.x,
... lambda y: F.cross_entropy(y, batch.target))
... opt.step()See Also
- lucid.compile.make_step—the cached entry point preferred for
repeated training loops — wraps this in a signature-keyedcache so the compile cost amortises.- lucid.compile.compile—module-level entrypoint.