compile(target: Callable[P, R], dynamic: bool = False)Wrap target so calls are routed through cached MPSGraph executables.
Accepts three call styles:
- Module wrapping —
lucid.compile(model)returns aCompiledModulethat delegates parameter walks /state_dict/ device moves tomodelwhile routing__call__through an executable cache. - Plain callable wrapping —
lucid.compile(fn)wherefnis a regular function or any callable. Returns a thin wrapper whose__call__traces the function's body on the first invocation with a new input signature and caches an executable. The wrapper has no parameters of its own — pure tensor-in / tensor-out. - Decorator usage —
@lucid.compileon a function (or@lucid.compile(dynamic=True)factory form) — identical to the plain-callable form above but lets the user opt-in inline.
Parameters
targetnn.Module or callableEither an
nn.Module instance (preferred when the
compiled unit carries learnable parameters) or any callable
whose signature is (*tensor_args) -> Tensor | tuple | dict.dynamicbool= FalseTreat the leading (batch) axis as symbolic so calls differing only
in batch size share one compiled executable (no recompile per
batch). Default
False. Convolutions cannot be symbolic-batch on
Apple's MPSGraph, so a conv graph transparently falls back to per-shape
static caching (any batch size still works, one compile per distinct
shape); non-conv graphs (MLP / Transformer / attention) get a single
dynamic executable.Returns
CompiledModuleA wrapper exposing the cache + run path. When target was a
Module the wrapper re-exposes the inner module's
parameters / state_dict / training mode; when it was a
plain callable those return empty sequences.
Examples
Module form:
compiled = lucid.compile(nn.Linear(8, 4).to('metal'))
y = compiled(x)
Callable form:
@lucid.compile
def f(x, y):
return (x @ y.T).relu()
z = f(a, b)
Factory decorator with kwargs:
@lucid.compile(dynamic=False)
def attention(q, k, v):
return F.scaled_dot_product_attention(q, k, v)See Also
CompiledModule—the wrapper class returned by this function.compile_optimizer—fused forward + backward + optimizer stepfor an even tighter training loop.